W HAT I S M ISSING IN IRM T RAINING AND E VALUATION ? C HALLENGES AND S OLUTIONS
A BSTRACT
Invariant risk minimization (IRM) has received increasing attention as a way to acquire environment-agnostic data representations and predictions, and as a principled solution for preventing spurious correlations from being learned and for improving models’ out-of-distribution generalization.
Yet, recent works have found that the optimality of the originally-proposed IRM optimization (IRM V 1) may be compromised in practice or could be impossible to achieve in some scenarios.
Therefore, a series of advanced IRM algorithms have been developed that show practical improvement over IRM V 1.
In this work, we revisit these recent
IRM advancements, and identify and resolve three practical limitations in IRM training and evaluation.
First, we find that the effect of batch size during training has been chronically overlooked in previous studies, leaving room for further improvement.
We propose small-batch training and highlight the improvements over a set of large-batch optimization techniques.
Second, we find that improper selection of evaluation environments could give a false sense of invariance for IRM.
To alleviate this effect, we leverage diversified test-time environments to precisely characterize the invariance of IRM when applied in practice.
Third, we revisit
Ahuja et al.
(2020)’s proposal to convert IRM into an ensemble game and identify a limitation when a single invariant predictor is desired instead of an ensemble of individual predictors.
We propose a new IRM variant to address this limitation based on a novel viewpoint of ensemble IRM games as consensus-constrained bilevel optimization.
Lastly, we conduct extensive experiments (covering 7 existing
IRM variants and 7 datasets) to justify the practical significance of revisiting IRM training and evaluation in a principled manner.
I NTRODUCTION
Deep neural networks (DNNs) have enjoyed unprecedented success in many real-world applications (He et al., 2016; Krizhevsky et al., 2017; Simonyan & Zisserman, 2014; Sun et al., 2014).
However, experimental evidence (Beery et al., 2018; De Haan et al., 2019; DeGrave et al., 2021; Geirhos et al., 2020; Zhang et al., 2022b) suggests that DNNs trained with empirical risk minimization ( ERM ), the most commonly used training method, are prone to reproducing spurious correlations in the training data (Beery et al., 2018; Sagawa et al., 2020).
This phenomenon causes performance degradation when facing distributional shifts at test time (Gulrajani & Lopez-Paz, 2020;
Koh et al., 2021;
Wang et al., 2022; Zhou et al., 2022a).
In response, the problem of invariant prediction arises to enforce the model trainer to learn stable and causal features (Beery et al., 2018; Sagawa et al., 2020).
In pursuit of out-of-distribution generalization, a new model training paradigm, termed invariant risk minimization ( IRM ) (Arjovsky et al., 2019), has received increasing attention to overcome the shortcomings of ERM against distribution shifts.
In contrast to ERM, IRM aims to learn a universal representation extractor, which can elicit an invariant predictor across multiple training environments.
However, different from ERM, the learning objective of IRM is highly non-trivial to optimize in practice.
Specifically, IRM requires solving a challenging bi-level optimization ( BLO ) problem with a hierarchical learning structure: invariant representation learning at the upper-level and invariant predictive modeling at the lower-level .
Various techniques have been developed to solve IRM effectively, such as (Ahuja et al., 2020;
Lin et al., 2022; Rame et al., 2022; Zhou et al., 2022b) to name a few.
Despite the proliferation of IRM advancements, several issues in the theory and practice have also appeared.
For example, recent works (Rosenfeld et al., 2020; Kamath et al.,
2021) revealed the theoretical failure of IRM in some cases.
In particular, there exist scenarios where the optimal invariant predictor is impossible to achieve, and the IRM performance may fall behind even that of ERM.
Practical studies also demonstrate that the performance of IRM rely on multiple factors, e.g. , model size (Lin et al., 2022;
Zhou et al., 2022b), environment difficulty (Dranker et al., 2021; Krueger et al., 2021), and dataset type (Gulrajani & Lopez-Paz, 2020).
Therefore, key challenges remain in deploying IRM to real-world applications.
In this work, we revisit recent IRM advancements and uncover and tackle several pitfalls in IRM training and evaluation, which have so far gone overlooked.
We first identify the large-batch training issue in existing
IRM algorithms, which prevents escape from bad local optima during IRM training.
Next, we show that evaluation of IRM performance with a single test-time environment could lead to an inaccurate assessment of prediction invariance, even if this test environment differs significantly from training environments.
Based on the above findings, we further develop a novel IRM variant, termed
BLOC -IRM, by interpreting and advancing the IRM-G AME method (Ahuja et al., 2020) through the lens of BLO with C onsensus prediction.
Below, we list our contributions ( ❶ - ❹ ).
❶ We demonstrate that the prevalent use of large-batch training leaves significant room for performance improvement in IRM, something chronically overlooked in the previous IRM studies with benchmark datasets C OLORED -MNIST and C OLORED -FMNIST.
By reviewing and comparing with 7 state-of-the-art (SOTA) IRM variants (Table 1), we show that simply using small-batch training improves generalization over a series of more involved large-batch optimization enhancements.
❷ We also show that an inappropriate evaluation metric could give a false sense of invariance to IRM.
Thus, we propose an extended evaluation scheme that quantifies both precision and ‘invariance’ across diverse testing environments.
❸ Further, we revisit and advance the IRM-G AME approach (Ahuja et al., 2020) through the lens of consensus-constrained BLO.
We remove the need for an ensemble (one per training environment) of predictors in IRM-G AME by proposing BLOC-IRM ( BLO with C onsensus IRM ), which produces a single invariant predictor.
❹ Lastly, we conduct extensive experiments (on 7 datasets, using diverse model architectures and training environments) to justify the practical significance of our findings and methods.
Notably, we conduct experiments on the C ELEB A dataset as a new IRM benchmark with realistic spurious correlations.
We show that BLOC-IRM outperforms all baselines in nearly all settings.
1. R ELATED W ORK
IRM methods.
Inspired by the invariance principle (Peters et al., 2016), Arjovsky et al.
(2019) define IRM as a BLO problem, and develop a relaxed single-level formulation, termed IRM V 1, for ease of training.
Recently, there has been considerable work to advance IRM techniques.
Examples of IRM variants include penalization on the variance of risks or loss gradients across training environments (Chang et al., 2020; Krueger et al., 2021; Rame et al., 2022; Xie et al., 2020; Xu &
Jaakkola, 2021; Xu et al., 2022), domain regret minimization (Jin et al., 2020), robust optimization over multiple domains (Xu & Jaakkola, 2021), sparsity-promoting invariant learning (Zhou et al., 2022b), Bayesian inference-baked IRM (Lin et al., 2022), and ensemble game over the environmentspecific predictors (Ahuja et al., 2020).
We refer readers to Section 2 and Table 1 for more details on the IRM methods that we will focus on in this work.
Despite the potential and popularity of IRM, some works have also shown the theoretical and practical limitations of current IRM algorithms.
Specifically, Chen et al.
; Kamath et al.
(2021) show that invariance learning via IRM could fail and be worse than ERM in some two-bit environment setups on C OLORED -MNIST, a synthetic benchmark dataset often used in IRM works.
The existence of failure cases of IRM is also theoretically shown by Rosenfeld et al.
(2020) for both linear and non-linear models.
Although subsequent IRM algorithms take these failure cases into account, there still exist huge gaps between theoretically desired IRM and its practical variants.
For example, Lin et al.
; Zhou et al.
(2022b) found many IRM variants incapable of maintaining graceful generalization on large and deep models.
Moreover, Ahuja et al.
; Dranker et al.
(2021) demonstrated that the performance of IRM algorithms could depend on practical details, e.g. , dataset size, sample efficiency, and environmental bias strength.
The above IRM limitations in- spire our work to study when and how we can turn the IRM advancements into effective solutions, to gain high-accuracy and stable invariant predictions in practical scenarios.
Domain generalization.
IRM is also closely related to domain generalization (Carlucci et al., 2019; Gulrajani & Lopez-Paz, 2020;
Koh et al., 2021;
Li et al., 2019;
Nam et al., 2021; Wang et al., 2022; Zhou et al., 2022a).
Compared to IRM, domain generalization includes a wider range of approaches to improve prediction accuracy against distributional shifts (Beery et al., 2018; Jean et al., 2016; Koh et al., 2021).
For example, an important line of research is to improve representation learning by encouraging cross-domain feature resemblance (Long et al., 2015; Tzeng et al., 2014).
The studies on domain generalization have also been conducted across different learning paradigms, e.g. , adversarial learning (Ganin et al., 2016), self-supervised learning (Carlucci et al., 2019), and meta-learning (Balaji et al., 2018;
Dou et al., 2019).
P RELIMINARIES AND S ETUP
In this section, we introduce the basics of IRM and provide an overview of our IRM case study.
IRM formulation.
In the original IRM framework Arjovsky et al.
(2019), consider a supervised learning paradigm, with datasets {D ( e ) } e ∈E tr collected from N training environments E tr = { 1 , 2 , . .
. , N } .
The training samples in D ( e ) (corresponding to the environment e ) are of the form ( x , y ) ∈
X × Y , where X and Y are, respectively, the raw feature space and the label space.
IRM aims to find an environment-agnostic data representation ϕ
θ : X → Z , which elicits an invariant prediction f w : Z → Y that is simultaneously optimal for all environments.
Here θ and w denote model parameters to be learned, and Z denotes the representation space.
Thus, IRM yields an invariant predictor f w ◦ ϕ θ : X → Y that can generalize to unseen test-time environments {D ( e ) } e/ ∈E tr .
Here ◦ denotes function composition, i.e. , f w ◦ ϕ θ ( · )
=
f w ( ϕ θ ( · )) .
We will use w ◦ θ as a shorthand for f w ◦ ϕ θ .
IRM constitutes the following BLO problem:
[Equation]
where ℓ ( e ) ( w ◦ θ ) is the per-environment training loss of the predictor w ◦ θ under D ( e ) .
Clearly, IRM involves two optimization levels that are coupled through the lower-level solution w ∗ ( θ ) .
Achieving the desired invariant prediction requires the solution sets of the individual lower-level problems { arg min ¯ w ℓ ( e ) ( ¯ w ◦ θ ) , e ∈ E tr } to be non-singleton.
However, BLO problems with non-singleton lower-level solution sets are significantly more challenging (Liu et al., 2021).
To circumvent this difficulty, Arjovsky et al.
(2019) relax (IRM) into a single-level optimization problem ( a.k.a. , IRMv1):
[Equation]
where γ > 0 is a regularization parameter and ∇ w | w =1 .
0 ℓ ( e ) denotes the gradient of ℓ ( e ) with respect to w , computed at w = 1 . 0 .
Compared with IRM, IRMv1 is restricted to linear invariant predictors, and penalizes the deviation of individual environment losses from stationarity to approach the lowerlevel optimality in (IRM).
IRMv1 uses the fact that a scalar predictor ( w = 1 . 0 ) is equivalent to a linear predictor.
Despite the practical simplicity of (IRMv1), it may fail to achieve the desired invariance (Chen et al., 2022; Kamath et al., 2021).
Case study of IRM methods.
As illustrated above, the objective of IRM is difficult to optimize, while IRMv1 only provides a sub-optimal solution.
Subsequent advances have attempted to reduce this gap.
In this work, we focus on 7 popular IRM variants and evaluate their invariant prediction performance over 7 datasets.
Table 1 and Table 2 respectively summarize the IRM methods and the datasets considered in this work.
We survey the most representative and effective IRM variants in the literature, which will also serve as our baselines in performance comparison.
Following Table 1, we first introduce the IRM V 0 variant, a generalization of IRMv1, by relaxing its assumption of linearity of the predictor w , yielding
[Equation]
Next, we consider the risk extrapolation method RE X (Krueger et al., 2021), an important baseline based on distributionally robust optimization for group shifts (Sagawa et al., 2019). Furthermore,
[Table] [Table] inspired by the empirical findings that the performance of IRM could be sensitive to model size (Choe et al., 2020; Gulrajani & Lopez-Paz, 2020), we choose the SOTA methods Bayesian IRM ( BIRM ) (Lin et al., 2022) and sparse IRM ( S PARSE IRM )
(Zhou et al., 2022b), both of which show improved performance with large models.
Also, we consider the SOTA method F ISHR
(Rame et al., 2022), which modifies IRM to penalize the domain-level gradient variance in single-level risk minimization.
F ISHR provably matches both domain-level risks and Hessians.
Lastly, we include IRM-G AME (Ahuja et al., 2020) as a special variant of IRM.
Different from the other methods which seek an invariant predictor, IRM-G AME endows each environment with a predictor, and leverages this ensemble of predictors to achieve invariant representation learning.
This is in contrast to other existing works which seek an invariant predictor.
Yet, we show in Section 5 that
IRM-G AME can be interpreted through the lens of consensus-constrained BLO and generalized for invariant prediction.
We also highlight that diverse dataset types are considered in this work (see Table 2) to benchmark IRM’s performance.
More details on dataset selections can be found in Appendix A.
L ARGE - BATCH T RAINING C HALLENGE AND I MPROVEMENT
In this section, we demonstrate and resolve the large-batch training challenge in current IRM implementations (Table 1).
[Figure]
Large-batch optimization causes instabilities of IRM training.
Using very large-size batches for model training can result in the model getting trapped near a bad local optima (Keskar et al., 2016).
This happens as a result of the lack of stochasticity in the training process, and is known to exist even in the ERM paradigm (Goyal et al., 2017; You et al., 2017a).
Yet, nearly all the existing IRM methods follow the training setup of IRM V 1 (Arjovsky et al., 2019), which used the full-batch gradient descent (GD) method rather than the mini-batch stochastic gradient descent (SGD) for IRM training over C OLORED -MNIST and C OLORED -FMNIST.
In the following, we show that large-batch training might give a false impression of the relative ranking of IRM performances.
We start with an exploration of the impact of batch size on the invariant prediction accuracy of existing IRM methods under C OLORED -MNIST.
Here the invariant prediction accuracy refers to the averaged accuracy of the invariant predictor applied to diverse test-time environments.
We defer its formal description to Section 4.
Figure 1 shows the invariant prediction accuracy of three IRM methods IRM V 1, IRM V 0, and RE X vs. the data batch size (see Figure A1 for results of other IRM variants and Figure A5 for C OLORED -FMNIST).
Recall that the full batch size (50k) was used in the existing IRM implementations (Arjovsky et al., 2019; Krueger et al., 2021).
As we can see, in the full-batch setup, IRM methods lead to widely different invariant prediction accuracies, where RE X and IRM V 1 significantly outperform IRM V 0.
In contrast, in the small-batch case (with size 1k), the discrepancy in accuracy across methods vanishes.
We see that IRMv0 can be as effective as IRMv1 and other IRM variants (such as RE X ) only if an appropriate small batch size is used.
Empirical evidence in Figure 1 shows that large-batch IRM training is less effective than smallbatch.
This is aligned with the observations in ERM (You et al., 2017b; 2018; 2019), where the lack of stochasticity makes the optimizer difficult to escape from a sharp local minimum.
We also justify this issue by visualizing the loss landscapes in Figure A2.
Notably, the small-batch training enables
IRM V 1 to converge to a local optimum with a flat loss landscape, indicating better generalization (Keskar et al., 2016).
Small-batch training is effective versus a zoo of large-batch optimization enhancements.
To mitigate the large-batch IRM training issue, we next investigate the effectiveness of both smallbatch training and a zoo of large-batch optimization enhancements.
Inspired by large-batch training techniques to scale up ERM, we consider Large-batch SGD ( LSGD )
(Goyal et al., 2017) and Layerwise Adaptive Learning Rate ( LALR )
(You et al., 2017b; 2018; 2019; Zhang et al., 2022a).
Both methods aim to smoothen the optimization trajectory by improving either the learning rate scheduler or the quality of initialization.
Furthermore, we adopt sharpness-aware minimization ( SAM ) (Foret et al., 2020) as another possible large-batch training solution to explicitly penalize the sharpness of the loss landscape.
We integrate the above optimization techniques with IRM, leading to the variants
IRM-LSGD , IRM-LALR , and IRM-SAM .
See Appendix B.1 for more details.
[Table]
In Table 3, we compare the performance of the simplest small-batch IRM training with that of those largebatch optimization technique-integrated IRM variants ( i.e. , ‘LSGD/LALR/SAM’ in the Table).
As we can see, the use of large-batch optimization techniques indeed improves the prediction accuracy over the original IRM implementation.
We also observe that the use of SAM for IRM is consistently better than LALR and LSGD, indicating the promise of SAM to scale up IRM with a large batch size.
Yet, the small-batch training protocol consistently outperforms large-batch training across all the IRM variants (see the column ‘Small’).
Additional experiment results in Section 6 show that smallbatch IRM training is effective across datasets, and promotes the invariance achieved by all methods.
M ULTI - ENVIRONMENT I NVARIANCE E VALUATION
In this section, we revisit the evaluation metric used in existing IRM methods, and show that expanding the diversity of test-time environments would improve the accuracy of invariance assessment.
Nearly all the existing IRM methods (including those listed in Table 1) follow the evaluation pipeline used in the vanilla IRM framework (Arjovsky et al., 2019), which assesses the performance of the learned invariant predictor on a single unseen test environment.
This test-time environment is significantly different from train-time environments.
For example, C OLORED -MNIST
(Arjovsky et al., 2019) suggests a principled way to define two-bit environments, widely-used for IRM dataset curation.
Specifically, the C OLORED -MNIST task is to predict the label of the handwritten digit groups (digits 0-4 for group 1 and digits 5-9 for group 2).
The digit number is also spuriously correlated with the digit color (Table 2).
This spurious correlation is controlled by an environment bias parameter β , which specifies different data environments with different levels of spurious correlation 1 .
In (Arjovsky et al., 2019), β = 0 .
1 and β = 0 .
2 are used to define two training environments , which sample the color ID by flipping the digit group label with probability 10% and 20% , respectively.
At test time , the invariant accuracy is evaluated on a single, unseen environment with β = 0 .
However, the prediction accuracy of IRM could be sensitive to the choice of test-time environment ( i.e. , the value of β ).
For the default test environment β = 0 .
9 , the predictor performance of three representative IRM methods (IRM V 1, IRM-G AME , F ISHR ) ranked from high to low is IRM-G AME > F ISHR > IRM V 1.
Given this apparent ranking, we explore more diverse test-time environments, generated by β ∈ Ω := { 0 . 05 , 0 . 1 , . . . [Figure]
Although the train- { 0 . 1 , 0 . 2 } belong to Ω , test data is generated afresh, different from training data.
We see in Figure 2A that the superiority of IRM-G AME at β = 0 .
9 vanishes for smaller β .
Consequently, for invariant prediction evaluated in other testing environments (e.g., β < 0 . 4 ), the performance ranking of the same methods becomes IRM V 1 > F ISHR > IRM-G AME .
This mismatch of results suggests we measure the ‘invariance’ of IRM methods against diverse test environments.
Otherwise, evaluation with single β could give a false sense of invariance.
In Figure 2B, we present the box plots of prediction accuracies for IRM variants, over the diverse set of testing environments ( β ∈ Ω) .
Evidently, IRM V 1, the oldest (sub-optimal) IRM method, yields the least variance of invariant prediction accuracies and the best average prediction accuracy, compared to both IRM-G AME and F ISHR .
To summarize, the new evaluation method, with diverse test environments, enables us to make a fair comparison of IRM methods implemented in different training environment settings.
Unless specified otherwise, we use the multi-environment evaluation method throughout this work.
A DVANCING IRM-G AME VIA C ONSENSUS - CONSTRAINED BLO
In this section, we revist and advance a special IRM variant, IRM-G AME (Ahuja et al., 2020), which endows each individual environment with a separate prediction head and converts IRM into an ensemble game over these multiple predictors.
Revisiting IRM-G AME .
We first introduce the setup of IRM-G AME following notations used in
Section 2.
The most essential difference between IRM-G AME and the vanilla IRM framework is that the former assigns each environment with an individual classifier w ( e ) , and then relies on the ensemble of these individual predictors, i.e. , 1 N (cid:80) e ∈E tr ( w ( e ) ◦ θ ) , for inference.
IRM-G AME is in a sharp contrast to IRM, where an environment-agnostic prediction head w ∗ simultaneously optimizes the losses across all environments.
Therefore, we raise the following question: Can IRM-G AME learn an invariant predictor?
Inspired by the above question, we explicitly enforce invariance by imposing a consensus prediction constraint C :
= { (cid:0) ¯ w (1) , ¯ w (2) , . . .
¯ w ( N ) (cid:1) | ¯ w (1)
= . . .
=
¯ w ( N ) } and integrate it with IRM-G AME .
Here, ¯ w ( e ) denotes the prediction head for the e -th environment.
Based on the newly-introduced constraint, the ensemble prediction head 1 N (cid:80) e ∈E tr w ( e ) can be interpreted as the average consensus over N environments: w ∗ :
= 1 N (cid:80) e ∈E tr w [Equation]
e ) = arg [Equation]
[Equation] e ) }
[Equation]
∈E tr
[Equation]
¯ w ( e ) − w ( e ) ∥ 22 .
With the above consensus interpretation, we can then cast the invariant predictor-baked IRM-G AME as a consensus-constrained BLO problem, extended from (IRM):
[Equation]
The above contains two lower-level problems: (I) per-environment risk minimization, and (II) projection onto the consensus constraint ( { w ( e ) } ∈ C ).
The incorporation of (II) is intended to ensure the use of invariant prediction head w ∗ ( θ ) in the upper-level optimization problem of (1).
Limitation of (1) and BLOC-IRM.
In (1), the introduced consensus-constrained lower-level problem might compromise the optimality of the lower-level solution w ∗ ( θ ) to the per-environment (unconstrained) risk minimization problem (I) , i.e. , violating the per-environment stationarity ∥∇ w ℓ [Equation] e ) ( w
[Equation]
( θ ) ◦ θ )
∥ 22 . Figure A3 justifies this side effect.
As we can see, the per-environment stationarity is hardly attained at the consensus prediction when solving (1).
This is not surprising since a constrained optimization solution might not be a stationary solution to minimizing the (unconstrained) objective function.
To alleviate this limitation, we improve (1) by explicitly promoting the per-environment stationarity ∥∇ w ℓ ( e ) ( w ∗ ( θ ) ◦ θ ) ∥ 22 in its upper-level problem through optimization over θ .
This leads to BLOC-IRM ( BLO with C onsensus IRM ): minimize
[Equation]
subject to Lower-level problems (I) and (II) in (1) [Equation] (BLOC-IRM) where γ > 0 is a regularization parameter like IRMv0.
Assisted by the (upper-level) prediction stationarity regularization, the consensus prediction (II) indeed simultaneously minimizes the risks of all the environments, supported by the empirical evidence that the convergence of ∥∇ w ℓ ( e ) ( w ∗ ( θ ) ◦ θ ) ∥ 22 towards 0 along each environment’s optimization path (see Figure A3).
Further, we elaborate on how the BLOC-IRM problem can be effectively solved using an ordinary
BLO solver.
First, it is worth noting that although both (IRM) and BLOC-IRM are BLO problems, the latter is easier to solve since the lower-level constraint (I) is unconstrained and separable over environments, and the consensus operation (II) is linear.
Based on these characteristics, the implicit gradient d w ∗ ( θ ) d θ can be directly computed as
[Equation]
[Figure]
Since the above lower-level problem is unconstrained , we can call the standard arg min differentiating method, such as implicit function approach (Gould et al., 2016) or gradient unrolling (Liu et al., 2021) to compute d w ( e ) ( θ ) d θ .
In our work, we adopt the gradient unrolling method, which approximates w ( e ) ( θ ) by a K -step gradient descent solution, noted by w ( e ) K ( θ ) and then leverages automatic differentiation (AD) to compute the derivative from w ( e ) K
( θ ) to the variable θ .
Figure 3 shows the working pipeline of BLOCIRM and its comparison to original IRM and
IRM-G AME methods.
We use K = 1 for the lower-level problem throughout our experiments.
We refer readers to Appendix B.2 for more algorithmic details.
We also explore the performance of our proposed BLOC-IRM with various regularization terms, based on the penalties used in the existing literature.
We show the best performance is always achieved when the stationarity is penalized in the upper-level (see Table A3).
E XPERIMENTS
In this section, we begin by introducing some key experiment setups (with details in Appendix C.1), and then empirically show the effectiveness of our proposed IRM training and evaluation improvements over existing IRM methods across various datasets, models, and learning environments.
6. E XPERIMENT S ETUPS
Datasets and models.
Our experiments are conducted over 7 datasets as referenced and shown in
Tables 1, 2.
Among these datasets, C OLORED -MNIST, C OLORED -FMNIST, C IFAR -MNIST, and
C OLORED -O BJECT are similarly curated, mimicking the pipeline of C OLORED -MNIST
(Arjovsky et al., 2019), by introducing an environment bias parameter ( e.g. , β for C OLORED -MNIST in Section 4) to customize the level of spurious correlation (as shown in Table 2) in different environments.
In the C ELEB A dataset, we choose the face attribute ‘smiling’ (vs. ‘non-smiling’) as the core feature aimed for classification, and regard another face attribute ‘hair color’ (‘blond’ vs. ‘dark’) as the source of spurious correlation imposed on the core feature.
By controlling the level of spurious correlation, we then create different training/testing environments in C ELEB A.
Furthermore, we study
PACS and VLCS datasets, which were used to benchmark domain generalization ability in the real world (Borlino et al., 2021).
It was recently shown by Gulrajani & Lopez-Paz (2020) that for these datasets, ERM could even be better than IRM V 1.
Yet, we will show that our proposed BLOC-IRM is a promising domain generalization method, which outperforms all the IRM baselines and ERM in practice.
In addition, we follow Arjovsky et al.
in adopting multi-layer perceptron (MLP) as the model for resolving C OLORED -MNIST and C OLORED -FMNIST problems.
In the other more complex datasets, we use the ResNet-18 architecture (He et al., 2016).
Baselines and implementation.
Our baselines include 7 IRM variants (Table 1) and ERM, which are implemented using their official repositories if available (see Appendix C.2).
Unless specified otherwise, our training pipeline uses the small-batch training setting.
By default, we use the batch size of 1024 for C OLORED -MNIST and C
OLORED
-FMNIST, and 256 for other datasets.
In Section 6.2 below, we also do a thorough comparison of large-batch vs small-batch IRM training. [Figure]
[Table]
Evaluation setup.
As proposed in Section 4, we use the multi-environment evaluation metric unless specified otherwise.
To capture both the accuracy and variance of invariant predictions across multiple testing environments, the average accuracy and the accuracy gap (the difference of the best-case and worst-case accuracy) are measured for IRM methods.
The resulting performance is reported in the form a ± b , with mean a and standard deviation b computed acrossindependent trials.
6. E XPERIMENT R ESULTS
Small-batch training improves all existing
IRM methods on C OLORED -MNIST &
C OLORED -FMNIST.
Recall from Sectionthat all the existing IRM methods (Table 1) adopt full-batch IRM training on C OLORED -MNIST & C OLORED -FMNIST, which raises the largebatch training problem.
In Table 4 , we conduct a thorough comparison between the originally-used full-batch IRM methods and their small-batch counterparts.
In addition, we present the performance of ERM and ERM-grayscale (we call it ‘grayscale’), where the latter is ERM on uncolored data.
In the absence of any spurious correlation in the training set, grayscale gives the best performance.
As discussed in Section 4 & 6.1, the
IRM performance is measured by the average accuracy and the accuracy gap across 19 testing environments, parameterized by the environment bias parameter β ∈ { 0 . 05 , . . .
We make some key observations from Table 4.
First , small batch size helps improve all the existing IRM methods consistently, evidenced by the 1% ∼ 3% improvement in average accuracy.
Second , the small-batch IRM training significantly reduces the variance of invariant predictions across different testing environments, evidenced by the decreased accuracy gap.
This implies that the small-batch IRM training can also help resolve the limitation of multi-environment evaluation for the existing IRM methods, like the sensitivity of IRM-G AME accuracy to β in Figure 2.
Third , we observe that IRM V 0, which does not seem to be useful in the large batch setting, becomes quite competitive with the other baselines in the small-batch setting.
Thus, large-batch could suppress the IRM performance for some methods.
In the rest of the experiments, we stick to the small-batch implementation of IRM training.
BLOC-IRM outperforms IRM baselines in various datasets.
Next, Table 5 demonstrates the effectiveness of our proposed BLOC-IRM approach versus ERM and existing IRM baselines across all the 7 datasets listed in Table 2.
Evidently, BLOC-IRM yields a higher average accuracy compared to all the baselines, together with the smallest accuracy gap in most cases.
Additionally, we observe that C ELEB A, PACS and VLCS are much more challenging datasets for capturing invariance through IRM, as evidenced by the small performance gap between ERM and IRM methods.
In
[Table]
particular, all the IRM methods, except F ISHR and BLOC-IRM, could even be worse than ERM on PACS and VLCS.
Here, we echo and extend the findings of Krueger et al.
(2021, Section 4.3).
However, we also show that BLOC-IRM is a quite competitive IRM variant when applied to realistic domain generalization datasets.
We also highlight that the C ELEB A experiment is newly constructed and performed in our work for invariance evaluation.
Like PACS and VLCS, this experiment also shows that ERM is a strong baseline, and among IRM-based methods, BLOC-IRM is the best-performing, both in terms of accuracy and variance of invariant predictions. [Figure]
[Figure]
IRM against model size and training environment variation.
Furthermore, we investigate the effect of model size and training environment diversity on the IRM performance.
The recent works (Lin et al., 2022; Zhou et al., 2022b) have empirically shown that IRM V 1 may suffer a significant performance loss when trained over large-sized neural network models, and thus developed BIRM and S PARSE IRM approaches as advancements of IRM V 1.
Inspired by these works, Figure 4 presents the sensitivity of invariant prediction to model size for different IRM methods on C OLORED -MNIST.
Here the model size is controlled by the dimension of the intermediate layer (denoted by d ) in MLP, and the default dimension is d = 390 ( i.e. , the vertical dotted line in Figure 4), which was used in (Arjovsky et al., 2019) and followed in the subsequent literature.
As we can see, when d > 390 , nearly all the studied IRM methods (including BLOC-IRM) suffer a performance drop.
Yet, as d ≥ 800 , from the perspective of prediction accuracy and model resilience together, the topbest IRM methods with model size resilience are BIRM, S PARSE IRM, and BLOC-IRM, although we did not intentionally design BLOC-IRM to resist performance degradation against model size.
We also show more experiment results in the Appendix. In Table A2 , we study IRM with different numbers of training environment configurations and observe the consistent improvement of BLOCIRM over other baselines.
In Table A4 we show that the performance of invariant prediction degrades, if additional covariate shifts (class, digit, and color imbalances on C OLORED -MNIST) are imposed on the training environments following Krueger et al.
(2021, Section 4.1) and also demonstrate that BLOC-IRM maintains the accuracy improvement over baselines with each variation.
In
Table A5 , we compare the performance of different methods in the failure cases of IRM pointed out by (Kamath et al., 2021) and show the consistent improvement brought by BLOC-IRM.
C ONCLUSION
In this work, we investigate existing IRM methods and reveal long-standing but chronically overlooked challenges involving IRM training and evaluation, which may lead to sub-optimal solutions and incomplete invariance assessment.
As a remedy, we propose small-batch training and multienvironment evaluation.
We reexamine the IRM-G AME method through the lens of consensusconstrained BLO, and develop a novel IRM variant, termed BLOC-IRM.
We conducted extensive experiments on 7 datasets and demonstrate that BLOC-IRM consistently improves all baselines. A CKNOWLEDGEMENT
The work of Y. Zhang and S. Liu was partially supported by National Science Foundation (NSF)
Grant IIS-2207052.
The work of M. Hong was supported by NSF grants CNS-2003033 and CIF1910385. The computing resources used in this work were partially supported by the MIT-IBM Watson AI Lab and the Institute for Cyber-Enabled Research (ICER) at Michigan State University. R EPRODUCIBILITY S TATEMENT
The authors have made an extensive effort to ensure the reproducibility of algorithms and results presented in the paper.
First, the details of the experiment settings have been elaborated in Section 6. and Appendix C.1.
In this paper, seven datasets are studied and the environment generation process for each dataset is described with details in Appendix A.
The evaluation metrics are also clearly introduced in Section 3.
Second, eight IRM-oriented methods (including our proposed BLOC-IRM) are studied in this work.
The implementation details of all the baseline methods are clearly presented in Appendix C.2, including the hyper-parameters tuning, model configuration, and used code bases.
For our proposed BLOC-IRM, we include all the implementation details in Section 5 and Appendix B.2, including training pipeline in Figure 3 and the pseudo-code in Algorithm A1.
Third , all the results are based on 10 independent trials with different random seeds.
The standard deviations are also reported to ensure fair comparisons across different methods.
Fourth, codes are available at https:/github.com/OPTML-Group/BLOC-IRM .
A PPENDIX A D ATASET S ELECTION
Compared to existing work, we expand the dataset types for evaluating the performance of different
IRM methods (see Table 2).
In addition to the most commonly-used benchmark datasets C OLOREDMNIST (Arjovsky et al., 2019) and C OLORED -FMNIST
(Ahuja et al., 2020), we also consider the datasets C IFAR -MNIST
(Lin et al., 2021; Shah et al., 2020) and C OLORED -O BJECT (Ahmed et al., 2020; Zhang et al., 2021), which impose artificial spurious correlations, MNIST digit number and object color, into the original CIFAR-10 and COCO Detection datasets, respectively.
Furthermore, we consider other three real-world datasets C ELEB A
(Liu et al., 2015), PACS (Li et al., 2017) and VLCS (Torralba & Efros, 2011), without imposing artificial spurious correlations.
Notably,
C ELEB A was first formalized and introduced to benchmark IRM performance.
The recent work (Gulrajani & Lopez-Paz, 2020) showed that when carefully implemented, ERM could outperform
IRM V 1 in PACS and VLCS.
Thus, we regard them as challenging datasets to capture invariance.
For C OLORED -O BJECT dataset, we strictly follow the setting adopted in (Lin et al., 2022) to generate the spurious features.
For C IFAR -MNIST we use the class “bird” and “plane” in the dataset
CIFAR as the invariant feature, while the digit “0” and “1” in MNIST as the spurious correlation.
C ELEB A dataset is, for the first time, introduced to measure IRM performance.
We select the attribute “Smiling” as the invariant label and use the attribute “Hair Color” (blond and black hair) to create a spurious correlation in each environment.
B I MPLEMENTATION D ETAILS B. D ETAILS ON L ARGE - BATCH O PTIMIZATION E NHANCEMENTS
✦ IRM-LSGD :
We first integrate large-batch SGD (LSGD) with IRM.
Following (Goyal et al., 2017), we make two main modifications: (1) scaling up learning rate linearly with batch size, and (2) prepending a warm-up optimization phase to IRM training.
We call the LSGD-baked IRM variant IRM-LSGD.
✦ IRM-LALR :
Next, we adopt layerwise adaptive learning rate (LALR) in IRM training.
Following (You et al., 2019), we advance the learning rate scheduler by assigning each layer of a neural network-based prediction model with an adaptive learning rate ( i.e. , proportional to the norm of updated model weights per layer).
More specifically, the model parameter update rule becomes:
[Equation]
where θ t,i denotes the i -th layer of the model parameters at iteration t , and u t,i represents the first-order gradient of the corresponding layer-wise model parameters.
We use τ ( ∥ θ t,i ∥ 22 = min { max {∥ θ t,i ∥ 22 , c l } , c u } ) as the scaling factor of the adaptive learning rate η t ∥ u t,i ∥ .
We use c l =and c u = 1 in our experiments.
✦ IRM-SAM : Lastly, we leverage sharpness-aware minimization (SAM) to simultaneously minimize the IRM loss and the loss sharpness.
The latter is achieved by explicitly penalizing the worstcase training loss of model weights when facing small weight perturbations.
This yields a wide minimum within a flat loss landscape.
More specifically, the sharpness-aware loss can be formulated as:
[Equation]
where
[Equation]
where the parameter perturbation ϵ is subject to the perturbation constraint ∥
ϵ
∥ 2 2 ≤ ρ .
When applied to IRM, we replace the per-environment training loss with the SAM loss, and adopt the ρ = 0 . 001 .
B. BLOC-IRM I MPLEMENTATION
As described in Section 5, the BLOC-IRM algorithm solves the IRM problem with two optimization levels.
We use 1 -step gradient descent to get the lower-level solution.
We retain the gradient graph in PyTorch to enable auto differentiation.
We assign each of the classification head { w ( e ) } a separate optimizer and use the same learning rate as the feature extractor θ .
For C OLORED -MNIST and C OLORED -FMNIST, we adopt a learning rate of 2 × 10 − 3 and use the Adam (Kingma & Ba, 2014) optimizer.
As for other datasets, we use the multi-step learning rate scheduler with an initial learning rate of 0 . 1 , which is consistent with other baselines.
We adopt the same penalty weight of 10 6 as IRM V 1 and IRM V 0.
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
e )
[Equation]
A3)
[Equation]
Upper
[Equation]
(A4) 6: end
C E XPERIMENTATION C. E NVIRONMENT S ETUP
As proposed in Section 4, we use the multi-environment evaluation metric unless specified otherwise.
To capture both the accuracy and variance of invariant predictions across multiple testing environments, the average accuracy and the accuracy gap (the difference between the best-case and worst-case accuracy) are evaluated for IRM methods.
Specifically, for the C OLORED -MNIST, C OLORED
-FMNIST, C OLORED -O BJECT , C IFARMNIST, and C ELEB A dataset, we manually create 19 test environments with uniformly sampled bias parameter β ∈ { 0 . 05 , 0 . 1 , . .
95 } , where the environment bias parameter β controls the spurious correlation (see Section 4 for more details).
For VLCS and PACS datasets, the training and test sets have 4 environments, namely { art painting, cartoon, sketch, photo } and { CALTECH, LABELME, PASCAL, SUN } respectively.
We use the first three environments as the training environments, while we use the test set of all four environments to form our proposed multi-environment invariance evaluation system.
C. B ASELINES
For each baseline method, we follow its official PyTorch repository except IRM-G AME and S PARSE IRM.
We translate the TensorFlow-based original code base of IRM-G AME to PyTorch.
As one of the latest IRM advancements, the official code of S PARSE IRM is not yet publicly available.
Therefore, we reproduce S PARSE IRM in PyTorch.
In particular, for C OLORED -MNIST and C OLORED -FMNIST, we stick to the original hyperparameters for the large-batch setting and tune the hyper-parameters of each method, including the penalty weight, number of warm-up epochs, and learning rate for the small batch setting.
In particular, for the large-batch setting, we use the penalty weight of 10 6 , 190 warm-up epochs, and 500 epochs in total, as suggested by the original IRM V 1 and inherited by its variants.
For the small-batch setting, we adopt the same penalty weight 10 6 .
Further, we found that the warm-up phase could be shortened without sacrificing accuracy.
Therefore, we use 50 warm-up epochs and total 200 epochs for all the methods.
For other datasets
[Figure]
we adopt the batch size of 128 and use ResNet-18 as the default model architecture.
We train for 200 epochs.
We adopt the step-wise learning rate scheduler with an initial learning rate of 0 . 1 .
The learning rate decays by 0 .
1 at the 100 th and 150 th epochs.
C. A DDITIONAL E XPERIMENT R ESULTS
The influence of batch size with all the baselines.
We show in Figure A1 the influence of training batch size on the performance of different methods.
We observe in Figure A1, as in Figure 1, that full batch setting does not achieve the best performance, and the use of mini-batch (stochastic gradient descent) indeed improves performance.
[Figure] [Figure]
Loss landscapes of IRM V 1 with different batch sizes.
We plot the loss landscapes of the models trained with IRM V 1 on C OLORED -MNIST using large (full) and small batch in Figure A2.
Using small batch training, IRM V 1 (Fig.
A2B) converges to a smooth neighborhood of a local optima.
This also corresponds to a flatter loss landscape than the landscape of the large-batch training (Figure A2(A)).
The loss landscapes demonstrate consistent results as other experiments discussed in
Section 3.
Training trajectory with BLOC-IRM with and without stationary loss.
In Figure A3, we plot the per-environment training trajectory of stationary loss when solving (1) and (BLOC-IRM) on
C OLORED -MNIST.
For (BLOC-IRM) we use the regularization term λ = 10 6 , which is aligned with the penalty coefficient used in IRM V 1.
As we can see, without the stationarity regularization, the stationary loss remains at a high level for both environments (the dotted curves).
Notably, the lower-level stationary can be reached fast with the stationarity penalty, as shown in the solid curves.
[Figure]
Performance of all the methods with full dataset list.
We show in Table A
the results of all the methods on the seven datasets we studied.
To be more specific, in Table A1, we append the results of
C OLORED -MNIST and C OLORED -FMNIST into Table 5 as a whole.
As we can see, our methods outperforms other baselines in all the datasets in terms of average accuracy, and stands top in most cases in terms of the accuracy gap.
[Figure]
[Table]
Experiment on different model sizes.
We show in Figure A4 the influence of the increasing model size on the performance of different baselines considered in this work.
Compared to Figure 4, we report additional standard deviation of the 10 independent trials in Figure A4. [Figure] [Equation]
[Figure]
Experiment with different training environments.
In Table A2, we show the performance of all the methods in more complex training environments, such as more training environments and more skewed environment bias parameter β .
As we can see, BLOC-IRM outperforms other baselines
[Table] [Table]
BLOC-IRM with different regularizations.
Based on the penalty terms used in the existing
IRM variants, we explore the performance of our proposed BLOC-IRM with various regularization, including the ones used in IRM V 1 (BLOC-IRM- V 1), RE X (BLOC-IRM- REX ), and F ISHR (BLOC-IRM-F ISHR ).
We conduct experiments on three different datasets and the results are shown in Table A3.
It is obvious that the best performance is always achieved when the per-environment stationarity is penalized in the upper-level.
This is not surprising since without an explicit promotion of stationarity, other forms of penalties do not guarantee the BLO algorithm to achieve an optimal solution.
[Table]
Performance comparison of different methods with additional covariate shifts.
Besides the sensitivity check on model size, Table A4 examines the resilience of IRM to variations in the training environment.
This study is motivated by Krueger et al.
(2021), who empirically showed that the performance of invariant prediction degrades if additional covariate shifts are imposed on the training environments.
Thus, we present the IRM performance on C OLORED -MNIST by introducing class, digit, and color imbalances, following Krueger et al.
(2021, Section 4.1).
Compared with
Table 4, IRM suffers a greater performance loss in Table A4, in the presence of training environment variations.
However, the proposed BLOC-IRM maintains the accuracy improvement over baselines with each variation.
In Table A2, we also study IRM with different numbers of training environments and observe the consistent improvement of BLOC-IRM over other baselines.
Exploration on the failure cases of previous IRM methods.
Some papers either theoretically (Rosenfeld et al., 2020) or empirically (Kamath et al., 2021) pointed out that the original IRM V 1 method could fail in certain circumstances, due to the fact that the regularization term used in IRM V 1 heavily relies on the “linear predictor” assumption.
Regarding this issue, we first bring to attention that the BLOC-IRM formulation does not require the predictors to be linear, since we adopt the regularization in the form of IRM V 0 in the upper-level objective, not IRM V 1.
To justify our argument, we repeat the experiments in (Kamath et al., 2021), which points out a specific scenario using the C OLORED -MNIST dataset where IRM V 1 fails.
[Table]
[Table]
More specifically, the models are trained in the training environments ( α, β ) =
(0 . 1 , 0 . 2) and (0 . 1 , 0 . 25) , and evaluated in the test environment (0 . 1 , 0 . 9) .
Note that denotes the label flipping rate and represents the environment bias parameter.
The results are shown in the Table A5.
As we can see, IRM V 1 is clearly worse than ERM as it achieves much lower average accuracy and higher accuracy gap.
However, BLOC-IRM outperforms ERM by obtaining high average accuracy and lower accuracy gap.
This result shows that BLOC-IRM seems promising to address the empirical IRM challenge discovered in (Kamath et al., 2021).
In the meantime, we also acknowledge that BLOC-IRM is not perfect since the advantage achieved by BLOC-IRM over ERM is not strong enough.
However, we stress that the main contribution of BLOC-IRM does not lie in solving the failure cases of IRM V 1, but to fix the issue of IRM-Game that resorts to a predictor ensemble to make the invariant prediction, which deviates from the spirit of acquiring invariant predictors in the original IRM paradigm.
A similar curve to Figure 1 on C OLORED -FMNIST.
We show the results for C OLOREDFMNIST similar to Figure 1 in Figure A5and
the conclusion does not change much.
As mentioned before, the large-batch training setup was typically used for IRM training over the C OLOREDMNIST and C OLORED -FMNIST datasets.
[Table]
[Figure]