Enhancing Unit-tests for Invariance Discovery
Abstract
Recently, Aubin et al.
(2021) proposed a set of linear low-dimensional problems to precisely evaluate different types of out-of-distribution generalization.
In this paper, we show that one of these problems can already be solved by established algorithms, simply by better hyper-parameter tuning.
We then propose an enhanced version of the linear unit-tests.
To the best of our hyperparameter search and within the set of algorithms evaluated, AND-mask is the best performing algorithm on this new suite of tests.
Our findings on synthetic data are further reinforced by experiments on an image classification task where we introduce spurious correlations.
1. Introduction
Over the last decade, deep learning has shown impressive performance in a variety of application domains, ranging from computer vision to natural language processing (Collobert & Weston, 2008; He et al., 2015).
However, recent years have been marked by a multitude of examples showing that deep learning models are prone to exploiting spurious correlations (Beery et al., 2018).
To address this issue,
AND-mask (Parascandolo et al., 2020) and several other works (Arjovsky et al., 2019; Koyama & Yamaguchi, 2020;
Khezeli et al., 2021) proposed to learn correlations that are invariant across multiple training distributions.
The key idea behind these methods is that when different interventions are used to produce the training datasets, invariant correlations should reflect the fixed causal mechanism underlying the target variable.
Aubin et al.
(2021) noted, however, that these algorithms perform poorly across a catalog of simple low-dimensional linear problems.
Since then, it has become common practice for a growing body of literature (Ahuja et al., 2021; Khezeli et al., 2021; Koyama & Yamaguchi, 2020; Du et al., 2021;
Wang et al., 2022; Nguyen et al.,
to evaluate the performance of algorithms for invariance discovery using the Linear Unit-Tests proposed by Aubin et al.
These unit-tests entail three classes of low-dimensional linear problems, each capturing a different structure for inducing spurious correlations.
Here, we will focus on the two linear classification tasks:
• Example 2 is a classification problem inspired by the cow vs. camel example (Beery et al., 2018) where spurious correlations are interpreted as background color.
• Example 3 is based on a classification experiment in Parascandolo et al.
(2020) where the spurious correlations provide a shortcut in minimizing the training error while the invariant classifier takes a more complex form.
In this paper we argue that, despite the conceptual appeal of these simple problems, evaluating an algorithm’s performance on these tasks can be misleading.
In Section 3, we show that Example 2 is a trivial problem that can be solved via empirical risk minimization (ERM) (Vapnik, 1998), while in Section 4 we show that Example 3 can become a much harder problem than intended.
Towards resolving these issues, we propose an enhanced version of the Linear Unit-Tests in Section 5.
We find that, to the best of our hyper-parameters search and within the set of algorithms evaluated, AND-mask (Parascandolo et al., 2020) is the only algorithm that effectively solves these new problems.
To support our findings on synthetic data, we evaluate
AND-mask on a more realistic image classification task in
Section 6 and find that it significantly outperforms ERM when spurious correlations are introduced.
2. Preliminaries
We consider n env environments; for each environment
[Equation] 1 we denote by D e = { x ei , y ei } n e i =
the corresponding dataset contaning n e samples.
The input feature vector x e = ( x e inv , x e spu ) ∈
R d contains features x e inv ∈ R d inv that elicits invariant correlations as well as features x
e spu ∈ R d spu that elicits spurious correlations.
Our goal is to construct invariant predictors that estimate y e by relying on x e inv and ignoring x e spu .
We provide here a brief overview of the algorithms compared in this work:
• Empirical Risk Minimization (Vapnik, 1998) minimizes the error on the union of all the training splits.
• Invariant Risk Minimization (Arjovsky et al., 2019) finds a representation of the features such that there exists a classifier, on top of that representation, that is simultaneously optimal for all environments.
• Inter-environmental Gradient Alignment (Koyama & Yamaguchi, 2020) minimizes the error on the training splits while reducing the variance of the gradient of the loss per environment.
• AND-mask (Parascandolo et al., 2020) minimizes the error on the training splits by updating the model on those directions where the sign of the gradient of the loss is the same for most environments.
• Information Bottleneck Invariant Risk Minimization
Ahuja et al., 2021) finds the predictor with the least entropy among all the highly predictive invariant predictors.
• Oracle (Aubin et al., 2021) is a version of ERM where features x e spu in the train set are shuffled at random across examples, hence spurious features are trivial to ignore
3. Example 2: Cows versus Camels
We consider Example 2 from the linear unit-tests of Aubin et al.
and prove that ERM, as well as all other algorithms considered in this paper, can solve this problem.
We report the data generating process here:
[Equation]
where 1 m ∈ R m denotes a vector of ones.
To construct the datasets D e for every e ∈ E we sample:
[Equation]
N
[Equation]
N
[Equation] y e ← (cid:40)if 1 ⊺ d inv x e inv > else , where for the first three environments the background probabilities are b e 0 = 0 .
95 , b e 1 = 0 .
97 , b e 2 = 0 .
99 and the animal probabilities are s e 0 = 0 .
3 , s e 1 = 0 .
5 , s e 2 = 0 .
When the number of environments are greater than three, then f e j ∼ Uniform (0 . 9 , 1) , and s e j ∼ Uniform (0 . 3 , 0 . 7) .
Proposition 1.
(Informal) Consider the setting of Example 2.
In this scenario, ERM, IRM, IB-IRM and ANDMask solve the OOD generalization problem.
The proof is a direct consequence of the results in (Ahuja et al., 2021) and can be found in the appendix. [Figure]
It should be noted that most of the experimental findings in the literature (Ahuja et al., 2021; Khezeli et al., 2021;
Koyama & Yamaguchi, 2020; Du et al., 2021; Wang et al., 2022; Nguyen et al., 2022) fail to achieve satisfactory performance in this example.
The issue lies with the hyperparameter search proposed in Aubin et al.
(2021), in particular, the learning rates tested are too small for convergence.
In light of Proposition 1, we argue that all the algorithms evaluated in the linear unit-tests should solve Example 2 .
We run experiments with higher learning rates and all the algorithms reach zero test error.
We report these results in
Figure 1, notably ANDMask converges significantly faster than all the other algorithms — even after tuning all hyperparameters (see Appendix A.1).
Moreover, ERM achieves zero error, making the task trivial for the other algorithms.
Overall, these results suggest that Example 2 , in its current form, is not effective for evaluating invariance discovery algorithms.
In Section 5, we address this issue and propose a more challenging version of the problem where an algorithm is ranked based on how much weight it puts on spurious features.
4. Example 3: Small Invariant Margin
We consider Example 3 from the linear unit-tests of Aubin et al.
This problem is a linear version of a classification experiment in Parascandolo et al.
(2020) where the spurious correlations provide a shortcut in minimizing the training error while the invariant classifier takes a more complex form.
Even though AND-mask was shown to solve the non-linear version of this problem in Parascandolo et al.
(2020), it has been experimentally observed that it fails to solve its linear counter-part.
In Section 4.2 we resolve this contradiction and provide an explanation for the failure of
AND-mask observed in Aubin et al.
; Ahuja et al.
] Following Aubin et al.
(2021), Example 3 is generated as follows: γ = 0 .
1 · 1 d inv , and µ e ∼ N d spu (0 , 1) , for all environments.
To construct the datasets D e for every e ∈ E we sample:
[Equation]
4.1. Can AND-mask solve Example 3?
We present a theoretical analysis of AND-mask in the setting of Example 3 , proving that it can solve the OOD generalization problem when the number of environments is sufficiently large.
For the sake of clarity, we consider the two dimensional case, i.e. d inv =
d spu = 1 .
We assume that the data is balanced and that AND-mask is initialized with weights w (0) =
(0 , 0) ⊤ — this choice can be interpreted as a prior belief that all features are spurious.
Along the lines of Parascandolo et al.
we define for every component [ m τ ] j of the mask m τ , [ m τ ] j = 1 (cid:104)
τ ⩽ 1 n env (cid:12)(cid:12)(cid:12)(cid:80) e sign (cid:16) [ ∇L e ( w )]
j (cid:17)(cid:12)(cid:12)(cid:12)(cid:105) , where ∇L e is the average gradient within an environment.
Our goal is to prove that AND-mask converges towards an invariant solution, that is a solution of the form:
[Equation]
First, we prove in Theorem 4.1 that the component of the mask m τ corresponding to the spurious feature is 0 with high probability when w = w (0) .
Theorem 4.1.
Let W =
( w 1 , w 2 ) ⊤ ∈ R 2 : w 2 = 0 } and [ m τ ] spu be the component of the mask corresponding to the spurious feature.
If w ∈ W then:
[Equation]
It follows that, if n env is big enough, [ w ] spu will likely never be updated during gradient descent, i.e
[Equation]
Next, we prove in Theorem 4.2 that the component of m τ corresponding to the invariant feature converges almost surely to 1 for all incorrect solutions of the form:
[Equation]
Theorem 4.2.
Let (cid:102) W ∗ be the set of incorrect solutions.
If w ∈ (cid:102) W ∗ then, as n e → ∞ , for all e ∈ E :
[Equation]
Now if w ( k ) ∈ W ∗ we are done.
If instead w ( k ) ∈ (cid:102) W ∗ then from Thm.
4.1 and Thm. 4.2, we know that gradient descent will only update the invariant feature and eventually converge towards a correct solution * .
We can conclude that with n env sufficiently large and proper initialization,
AND-mask will only rely on invariant features and converge towards the invariant solution of Example 3 .
4.2. A Tough Problem
Our theoretical analysis suggests that AND-mask should correctly identify invariant signals in the setting of Example 3 and yet experimental evidence is contradictory (Aubin et al., 2021; Ahuja et al., 2021).
We argue that the problem of invariance discovery becomes much harder when constructing the dataset from a small number of environments.
Proposition 2.
Consider the setting of Example 3.
We have that the sign of at least one component [ µ e ] j will be the same in all the environments with probability p = 1 − ( [Equation] − 2 1 − n env ) d spu .
Note that, when sign([ µ e ] j ) is constant over all the environments it becomes much hard to distinguish spurious signals from invariant signals.
Now, if we consider the setting ( d inv , d spu , n env ) = (5 , 5 , 3) adopted in Ahuja et al.
Aubin et al.
(2021), with probability p ≈ 0 .
97 at least one spurious dimension will have constant sign across environments, hence explained why no algorithm could solve the problem with these settings.
In Section 5, we propose a modified version of Example 3 where we force the sign of spurious signals to be inconsistent in at least one environment.
5. Enhancing Linear Unit-Tests
One way to make Example 2 more challenging would be to change the support of the spurious features distribution at test-time.
In that case, both ERM and IRM fail to solve the OOD generalization problem as proved in Thm.
3 (Insufficiency) of Ahuja et al.
However, changing the support of the spurious features does not provide a provable guarantee of robustess; instead, we propose to rank algorithms based on how much weight they place on spurious features.
We denote by w spu the components of the weight vector corresponding to the spurious features.
Note that, if ∥ w spu ∥ ∞ > 0 we can always find a distribution shift that would make the algorithm fail.
We consider the task to be solved if ∥ w spu ∥ ∞ = 0 .
We run experiments in this new scenario and report the results in Figure 3.
To the best of our hyper-parameters search, AND-mask is the only algorithm that solves this [Figure] example correctly, IB-IRM comes close but it could still fail for some significant distribution shift of the spurious features.
The remaining algorithms fail to solve this task.
[Figure]
[Figure]
IRMv
IGA
[Figure]
IRM
[Figure] As for Example we propose a modified version of the problem to complement the more challenging one proposed by Aubin et al.
In particular, to mitigate the issue raised in Prop.
we force the sign of the spurious mean in the last environment to be the opposite of the sign of the spurious mean in the first environment, i e. sign( µ n env ) = − sign( µ 1 ) .
This simple modification introduces an asymmetry in the dataset which makes it possible for an algorithm to identify, in principle, the invariant features — in contrast to the data generation process outlined in
Section 4.2.
We run experiments for this modified version of
Example 3 and report the results in Table 1.
We observe that, to the best of our hyper-parameters search, only AND-mask can effectively solve the OOD generalization problem. appendix.
[Table]
6. Image Classification Task
To further assess AND-mask’s performance under more realistic conditions, we construct a binary classification problem for images with spurious correlations.
The task is to correctly classify dog vs. cat; at training-time each class is associated with n spurious objects and at test-time these associations are reversed.
In Figure 2 we show a few examples of the dataset for n = 3 .
The main challenge is learning a classifier that does not rely on spurious features, i [Table] e. the associated objects [Table]
and thus can generalize well in the test environment.
We extract supervised representations from the last layer of BiT-M-R152 (Kolesnikov et al., 2020).
Then, we train a linear classifier on top of these representations using both AND-mask and ERM.
We report the results of this experiment in Table 2.
Overall, AND-mask consistently outperforms ERM — it is worthwhile to note how the gap in performance between AND-mask and ERM increases once spurious associations are introduced
[Table] [Table]
7. Conclusions
Out of distribution generalization is one of the most challenging issues in machine learning, and properly evaluating an algorithm’s performance in this context is essential.
This paper extends the original set of problems proposed by
Aubin et al.
(2021) with the purpose of providing a more fine-grained evaluation of invariance discovery algorithms.
We hope researchers will continue to extend this suite of tests to learn more about the strengths and weaknesses of new algorithms in a transparent and standardized manner.
A. Appendix A.1. Experiments details
We used the published code of Arjovsky et al.
(2019) to conduct all experiments.
We follow the same protocol as prescribed in Aubin et al.
(2021) for model selection, hyper-parameter selection, training, and evaluation.
For all the examples, the models used are linear.
The training loss is the binary cross-entropy for the classification setting.
For AND-mask we use
Adam optimizer while for all the other algorithms we use GD without any batching.
Hyper-parameter tuning for Figure 1 and Figure 3 .
We run 100 hyper-parameter queries for each model with 5 data seeds.
AND-mask treats every data point as coming from its own environment and we initialize all the weights of the
AND-mask network at 0 † .
Below we report how the search over the hyper-parameter is performed.
[Table]
EARNING
[Table]
ATE
W EIGHT D
[Table]
λ IRM λ IB τ
IGA P ENALTY
ERM 10 Uniform( − 2 , − 1) 10 Uniform( − 6 , − 2) ✗ ✗ ✗ ✗
IRM V10 Uniform( − 2 , − 1) 10 Uniform( − 6 , − 2) 1 − 10 Uniform( − 3 , − 0 .
3) ✗ ✗ ✗
IB-IRM 10 Uniform( − 2 , − 1) 10 Uniform( − 6 , − 2) 1 − 10 Uniform( − 3 , − 0 . 3)
1 − 10 Uniform( − 2 , 0) ✗ ✗
AND- MASK 10 Uniform( − 2 , − 1)
10 Uniform( − 6 , − 2) ✗ ✗
Uniform(0 . 8 , 1 . 0) ✗
IGA 10 Uniform( − 2 , − 1) 10 Uniform( − 6 , − 2)
✗ ✗ ✗ 10 Uniform(1 , 5)
Hyper-parameter tuning for Table 1 .
We run 150 hyper-parameter queries and average over 3 data seeds.
Below we report how the search over the hyper-parameter is performed.
[Table]
W
[Table]
λ IRM λ
[Table]
τ
IGA P ENALTY
ERM
[Table]
[Table]
) ✗ ✗ ✗ ✗
IRM V
[Table] [Table]
[Table]
✗ ✗ ✗
IB-IRM
[Table] [Table]
[Table]
✗ ✗
AND- MASK
[Table]
✗ ✗
[Table]
✗
IGA
[Table]
✗ ✗
[Table]
Image Classification Task .
Note that there is no standard notion of environments here, which is why AND-mask treats every example as coming from its own environment.
This assumption is not unreasonable, as every image in the dataset was literally collected in a different physical environment.
We train both AND-mask and ERM using GD without any batching.
A.2. Example 2 Proofs
We provide here a proof for Proposition 1.
Proof.
Example 2 follows the linear classification SEM (FIIF) with zero noise from Assumption 2 in Ahuja et al.
(2021).The invariant features are strictly separable, bounded and satisfy support overlap; the spurious features are bounded and satisfy support overlap.
It follows from Theorem 3 (Sufficiency) in Ahuja et al.
(2021) that both ERM and IRM solve the OOD generalization problem.
As ERM and IRM are edge cases respectively of ANDMask and IB-IRM, for τ = 0 and γ = 0 , it follows that they can both solve the OOD generalization problem as well.
A.3. Example 3 Proofs
We consider datasets {D e } e ∈E , and D e
=
( x e i , y e i ) , i e = 1 , . . . , n e .
Here x e i ∈
X ⊆
R m is the vector containing the observed inputs,
w ∈ R
m is the vector of weights and y ei ∈ { 0 , 1 } are the targets.
The superscript e ∈ E can be interpreted as an environment label.
We define α w ∈ R d as a (signed) measure of gradient agreement across environments [ α w ] i := 1 |E| (cid:80) e ∈E sign([ ∇L e ( w )] i ) .
We report here the data generating process of Example 3.
Let γ = 0 .
1 · 1 d inv , and µ e ∼ N d spu (0 , 1) , for all environments.
To construct the datasets D e for every e ∈ E we sample:
[Equation]
We analyse learning by minimizing an empirical loss of the form:
[Equation]
The following lemma characterizes the asymptotic behaviour of ∇L e as n e → ∞ :
Lemma A.1.
Let P ( θ 1 ,θ 2 ) be the distribution of N (cid:0) [ θ 1 , θ 2 ] T , Iσ 2 (cid:1) .
If {D e } e ∈E are sampled according to Example 3 then the following holds:
[Equation]
Proof.
It suffices to apply the law of large numbers,
i
[Equation]
e. sums converge to expectations:
[Equation]
This concludes the proof.
We now present the proof of Theorem 4.1.
Proof.
We will denote with [ ∇L e ( w )] spu the components of the gradient corresponding to the spurious features.
Similarly [ x ] inv and [ x ] spu will denote the components of the feature vector correspoding to invariant and spurious features, respectively.
From assumptions we have w =
( k, 0) ⊤ for some k ∈ R .
Applying Lemma A.1 we have that:
[Equation]
Next, the fact that σ ( x ) = 1 − σ ( − x ) and x ∼ P − γ = ⇒ − x ∼ P γ , we get
[Equation]
It follows that sign (cid:16) [ ∇L e ( w )] spu (cid:17) does not depend on w :
[Equation]
Moreover, we have that for spurious signals α w converges to:
[Equation]
Now, recall that µ e ∼ N (0 , 1) are independent and normally distributed.
Therefore sign( µ 1 ) , . .
, sign( µ n env ) are independent bounded random variables with sign( µ i ) ∈ {− 1 , 1 } for all i .
We can bound this sum using Hoeffding’s inequality:
[Equation]
and it follows that:
[Equation]
Next, we provide a proof for Theorem 4.2.
Proof.
First, we prove the statement for w 0 = (0 , 0)
⊤ .
Applying Lemma A.1 we have that:
[Equation]
Then it follows that:
[Equation]
Now, to prove that [ α w ] inv a.s. −→ 1 for all w ∈ (cid:102) W ∗ = { ( w 1 , w 2 ) ⊤ ∈ R 2 : w 2 =
∧ w 1 ≥ 0 } it suffices to show that [ ∇L e ( w 0 )] inv ≤ [ ∇L e ( w )] inv for all w ∈ (cid:102) W ∗ .
It is enough to prove that the first derivative of (cid:2) ∇L e (( w,
[Table]
positive
[Table]
[Equation]
Finally, to conclude the proof we have:
[Equation]
m
[Equation]
Finally, we provide the proof of Proposition
[Equation]
Proof
[Equation] [Equation]