What Is Missing in IRM Training and Evaluation? Challenges and Solutions
Authors: Yihua Zhang
Pranay Sharma
Parikshit Ram
Mingyi Hong
Kush R. Varshney
Sijia Liu
Venue: ICLR 2023 poster (2023)
Keywords: invariant risk minimization
bi-level optimization

Abstract: Invariant risk minimization (IRM) has received increasing attention as a way to acquire environment-agnostic data representations and predictions, and also a principled solution for preventing spurious correlations from being learned and improving models’ out-of-distribution generalization. Yet, recent works have found that the optimality of the originally-proposed IRM optimization (IRMV1) 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 IRMV1. 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 bi-level 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.

Published as a conference paper at ICLR 2023
WHAT Is MISSING IN IRM TRAINING AND EVALUA- TION? CHALLENGES AND SOLUTIONS
Yihua Zhang', Pranay Sharma", Parikshit Ram*, Mingyi Hong', Kush Varshney', Sijia Liu Michigan State University, CCarnegie Mellon University, SIBM Research, *University of Minnesota
ABSTRACT
Invariant risk minimization (IRM) has received increasing attention as a way to acquire petironment-agnoosee data representations and predictions, and as a prin- cipled solution for preventing spurious correlations from being learned and for im- proving models' out-of-distribution generalization. Yet, recent works have found that the optimality of the originally-proppsee IRM optimization (IRMV1) may be compromised in practice could be impossible to achieve in some scenar- ios. Therefore, a series of advanced IRM algorithms have been developed that show practical improvement over IRMvI. In this work, 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 im- provement. We propose small-batch training and highlight the improvements over a set of large-batch optimization techniques. Second, we find that improper se- lection 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 Ahujaa et al. (2020)'s proposal to convert IRM into ensemble game and identify a limitation when a single invariant predictor is desired instead of an ensemble individual predictors. We propose based on a novel viewpoint of ensemble IRM games as consensus consensus-constraing constrained bi- level optimization. Lastly, IRM variants and 7 datasets) to justify the practical significance of revisiting IRM training and evaluation in principled manner.
IRM variant address this limitation
conduct extensive experiments (covering 7 existing
INTRODUCTION
Deep neural networks (DNNs) have enjoyed unprecedented success in many real-world applications (He et al.] 2016 Krizhevsky al. 2017 Simonyan & Zisserman 2014 Sun et al.) 2014). However, experimental evidence (Beery al. 2018 De Haan et al. 2219 DeGrave et 2022 | Geirhos et al.) 2020 Zhang et al.] (022b) suggests that DNNs trained with empirical risk minimization (ERM), the prone 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 cet al. 2022 Zhou et al. 2022a). In response, the problem of invariant prediction arises to enforce
most commonly used training method,
the model trainer learn stable and causal features (Beery et 2018 Sagawa et al. 2020).
In pursuit of out-of-distributie 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 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 (Ahuja et al. 2020; Lin al. 2022 Rame al. 2022 Zhou et al. 20226) to name a few. Despite the proliferation IRM advancements, several issues in the theory and practice have also appeared. For example, recent works (Rosenfeld et al. 2020 Kamath al.

---

Published as a conference paper at ICLR 2023
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 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 evalu- ation, 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 train- ing environments. Based on the above findings, we further develop novel IRM variant, termed BLOC-IRM, by interpreting and advancing the IRM- method (Ahuja et al.) 2020 through the lens of BLO with Consensus prediction. Below, we list our contributions (0-0).
We demonstrate that the prevalent use of large-batch training leaves significant room for perfor- mance improvement in IRM, something chronically overlooked in the previous IRM studies with benchmark datasets COLORED- DOLISTD-MNIST and COLORED-FMNIST. By reviewing and comparing with 7 state of-thee-rttttttoo (SOTA) IRM variants [T, we show that simply using small bbtl-b train- ing improves generalization over a series of more involved large-batch optimization enhancements.
@ We also show that inappropriate evaluation metric could give 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-GAME approach (Ahuja et al. 2020) through the lens of consensus-constraiined BLO. We remove the need for ensemble (one per training environment) of predictors in IRM-GAME by proposing BLOC-IRM (BLO with Consensus IRM), which pro- duces 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 CELEBA dataset as a new IRM benchmark with realistic spurious correlations. We show that BLOC-IRM outperforms all baselines in nearly all settings.
1.1 RELATED WORK
IRM methods. Inspired by the invariance principle (Peters al.) (016), Arjovsky et al. (2019) define IRM BLO problem, and develop relaxed single-level formulation, termed IRMv] 1, for
ease of training. Recently, there has been considerable work to advance IRM techniques. Exam-
ples of IRM variants include penalization on the variance of risks loss gradients across training environments (Chang et al. [2020] Krueger et al.] 2021 Rame et al. [002]] Xie et al. 22000: Xu & Jaakkola | al. (022)) domain regret minimization (Jin al. (2020), robust optimization over multiple domains (Xu & Jaakkola 2021),
invariant learning (Zhou et al.)
2022b), Bayesian inference-baked IRM (Lin et al. (2022). and ensemble game over the environment- specific predictors (Ahuja et al. 2020). We refer readers to Section and Table 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 prac- tical limitations of current IRM algorithms. Specifically, Chen al. (2022); Kamath al. (202T) show that invariance learning via IRM could fail and be worse than ERM in some two-bit environ- ment setups on COLORED-MNISS 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. 2021 2022); Zhou et al. (2022b) found many IRM variants incapable of main- taining graceful generalization on large and deep models. Moreover, Ahuja et al. (022); } Dranker cet al. (02[) 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-
a synthetic benchmark dataset often used in IRM works. The
2

---

Published as a conference paper at ICLR 2023
spire our work study when and how we can turn the IRM advancements into effective solutions, to gain high- righ-accuracy and stable invariant predictions in practical scenarios.
Domain generalization. IRM 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
al. 2022 Zhou al. 2022a). Compared to IRM, domain generalization includes wider range
of approaches to improve prediction accuracy against distributional shifts (Beery et al.
2018 0 Jean
al. 2016 et (02I). For example, an important line of research is to improve representation
learning by encouraging cross 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)), sulf-supervised learning (Carlucei et al. 2019). and meta-learning (Balaji et al. 2018 Dou et al. 2019).
2 PRELIMINARIES AND SETUP
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 super- vised learning paradigm, with datasets (D(e) Jeee.. collected from N training environments Eur = 21,2.... ..... N]. The training samples in D(e) (corresponding to the environment e) are of the form (x. y) E X x V, where X and Y are, respectively, the raw feature space and the label space. IRM aims to find an environment-quee angnoettt-- data representation òe : X + Z, which elicits an invariant + y that is simultaneously optimal for all environments. Here O and w denote model parameters to be learned, and 2 denotes the representation space. Thus, IRM yields an invari- ani predictor fw the X + y that can generalize to unseen test-time environments (D(c)) JeeE.. Here o denotes function composition, i.e., fw (.)(.) = (whe We will use o as a shorthand for fw Φg. IRM constitutes the following BLO problem:
prediction fw
minimize Eeer,
glo) (0)
subject to w'(0) E arg min gl0) (w 0), Ve Ex,,
(IRM)
where g(t) (wol) is the per-environment training loss of the predictor wef under Dle). Clearly, IRM involves two optimization levels that are coupled through the lower-level solution (0). Achiev- ing the desired invariant prediction requires the solution sets of the individual lower-level problems (arg min, g(e) 0), } 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 dif- ficulty, Arjovsky et al. (2019) relax (RMM into a singe--level optimization problem (a.k.a,, IRMv1):
minimize
Ecee, [ele) (0) T wol (w 0)])
(IRMv1)
where y >> 0 is a regularization parameter and V kdlw-100 ( gle) denotes the gradient of g(e) with respect w, computed
= 1.0. Compared with IRM IRMvi is restricted to linear invariant predictors, and penalizes the deviation of individual environment losses from stationarity to approach the lower- level optimality in (IRM). [RMvl] uses the fact that a scalar predictor
1.0) is equivalent to a linear predictor. Despite the practical simplicity of [RMvv]. 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 [RM is difficult to optimize, while IRMvT 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 and Tablee]rsppectivll 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 I first introduce the IRMv0 variant, generalization of [RMvl by relaxing its assumption of linearity of the predictor w, yielding
minimize w
1e (w e 0) + > ele) (le) (w 0) [B].
(IRMv0)
Next, we consider the risk extrapolation method REx (Kueeer et al. 2021). an important baseline based distributionally robust optimization for group shifts (Sagawa al. 2019). Furthermore,
3

---

Published as a conference paper at ICLR 2023
Table I: Summary of the 7 existing IRM variants consid- ered in this work, and the proposed BLOC-IRM method (see Scction We also list the 7 benchmark datasets used to evaluate IRM performance. namely. COLORED-MNIST (CoM), FOLNIS (CoF), CIFAR-MNIST (CiM), LLLORED-OBJECE (CoO), CELEBA (CA), PACS (P) and VLCS (A). The symbols ' signifies the datasct used the specific reference.
Table 2: Dataset setups. 'Invariant' and 'Spurious' represent the core and spurious fcatures. 'Envl" and "En22" are environ- ments with different spurious correlations.
Dataset
Invariant
Spurious
Env
Eas 2
CoM
Digit
Color
CoF
Object
Color
Metbod
IRM
CoM
Reference
GM
CIFAR
MNIST
IRMvn IRM-GAME
Arjonscy
CoO
Object
Color
ICMI
CA
Smiling
Hhir Color
BIRM
CSPR
SPARSEIRM
JEMI
Object
Texture
PUSHER
ICMI
Outs
NIA
This Work
v Object
Environment
inspired by the empirical findings that the performance of IRM could be sensitive to model size (Choe et al. 2020 Gulrajani & Lopez-Par (020), we choose the SOTA methods Bayesian IRM (BIRM) (Lin al. 2022) and sparse IRM (SPARSEIRM) (Zhou et al. 2022b), both of which show improved performance with large models. Also, we consider the SOTA method FISHR (Rame ctt al., 2022), which modifies IRM to penalize the domain-level gradient variance in single-level risk minimization. FISHR provably matches both demain-level risks and Hessians. Lastly, we include IRM-GAME (Ahuja et al. (020))
special variant of IRM. Different from the other methods which seek an invariant predictor, IRM-GAME endows each environment with a predictor, and leverages this ensemble of predictors achieve invariant representation learning. This is in contrast to other existing works which seek an invariant predictor. Yet, we show in Section that IR.M-AME can be interpreted through the lens of consensus-cowd consttrined 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
3 LARGE-BATCH TRAINING CHALLENGE AND IMPROVEMENT
In this section, we demonstrate and resolve the large-batch training challenge in current IRM imple- mentations (Table
Large-batch optimization causes instabilities of IRM train-
70
ing. Using very large-size batches for model training can
Invariant Accuracy (%)
69
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 ist even in the ERM paradigm (Goyal et al.) [2117] You et al.] 2017a). Yet, nearly all the existing IRM methods follow the training setup of IRMv1 (Arjovsky et al. 2019), which used the full-batch gradient descent (GD) method rather than the mini-batch stochastic gradient descent (SGD) for IRM train- ing over COLORED-MNIST and COLORTD-FMNIST. In the following, we show that large-batch training might give a false impression of the relative ranking of IRM performances.
68
67
IRM-V1
66
IRM-V0
65
REx
64
256
IK 2k 4k 50k
Batch Size
Figure 1: The performance of three IRM methods (IRMv1, IRMv0, and REx) vs. batch-size under COLORED- MNIST. The full batch-size is 50k.
We start with an exploration of the impact of batch size on the invariant prediction accuracy of existing IRM methods under COLORED-MMN -NIST. 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 Figure [ shows the invariant prediction accuracy of three IRM methods IRMv1. I, IRMv0, and REX vs. the data batch size (see Tigure|A]] for results of other IRM variants and Figure A5 for MILLORDDFMMIISE the existing IRM implementations (Arjovsky et al. 2019: Krueger et al. 2021). As can see, in the full- batch setup, IRM methods lead to widely different invariant prediction accuracies, where REx and IRMv1 significantly outperform IRMvO. In contrast, in the small-batch case (with size
Recall that the full batch size (50k) was used in
4

---

Published as a conference paper at ICLR 2023
1k), the discrepancy in accuracy across methods vanishes. We see that [RRMO can be as effective as IRMvI and other IRM variants (such as REx) only if an appropriate small batch size is used.
Empirical evidence in Figure[] shows that large-batch IRM training is less effective than small- batch. 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 sharp local minimum. We also justify this issue by visualizing the loss landscapes in Figure[A2 Notably, the small-batch training enables
IRMVI to converge to a local optimum with a flat loss landscape, indicating better generalization
Kessar et al. 2016).
Small-batch training is effective versus zoo of large-batch optimization enhancements. To mitigate the large-batch IRM training issue, we next investigate the effectiveness of both small- batch training and a zoo of large-batch optimization enhancements. Inspired by large-batch training techniques to scale up ERM, we consider Large- SGD (LSGD) (Goyal et al. 2017) and Layer- wise Adaptive Learning Rate (LALR) YYou 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-awe -waare minimization (SAM) (Foret et al. (2000) 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.I for more details.
In Table we compare the performance of the sim- plest small-batch IRM training with that of those large- batch optimization technique-integrann (i.e., 'LSGD/LALR/SA 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 proto- col consistently outperforms large-batch training across all the IRM variants (see the column 'Smmall'.. Addi- tional experiment results in Section show that small- batch IRM training is effective across datasets, and promotes the invariance achieved by all methods.
Table 3:
Prediction accuracy IRM methods on COLORED-MNIST using the original large-batch implementation (" ('Origi- nal'), the large-batch optimization triimization-ntegraaedd implementations ('LSGDILALR/SAM) and the small batch training recipe ('Small').
IRM variants in the Table). As we can see,
Method
Original
LSGD
LALR
SAM
Small
IRMV1
67.13
67.31
67.44
67.79
68.33
IRMVO
65.39
66.42
66.76
66.99
68.37
IRM-GAME
67.73
REX
67.42
67.53
67.59
67.82
68.42
BIRM
67.93
67.99
68:21
6M.32
68.71
SPARSHIM
67.72
67.85
67.99
68.13
68.81
67.93
Average
67.02
67.25
67.34
67.63
68.44
4 MULTI-ENVIRONMENT INVARIANCE EVALUATION
In this section, we revisit the evaluation metric used in existing IRM methods, and show that expand- ing the diversity of test-time environments would improve the accuracy of invariance assessment.
Nearly all the existing IRM methods (including those listed in Table follow the evaluation pipeline used in the vanilla IRM framework (Arjovsky et al.. 2019). which assesses the performance of the learned invariant predictor on single unseen test environment. This test-time environment is signif- icantly different from train-time environments. For example, COLORED-MNIST (Arjovsky et al.) 2019) suggests a principled way to define two-bit environments, widely-used for IRM dataset cura- tion. Specifically, the MOLLRRDD-MISS' 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 This spurious correlation is controlled by an environment bias param- eler 3, which specifies different data environments with different levels of spurious correlation "I In (Arjovsky al. 2019), 3 = 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 3 = 0.9.
However, the prediction accuracy of IRM could be sensitive to the choice of test-time en- vironment (i.e., the value of 8). For the default test environment 8 0.9, the predic- tor performance of three representative IRM methods (IRMv1, IRM-GAME, FISHR) ranked Given this apparent ranking, we ex-
from high to low is plore more diverse test-time environments, generated by 3 E Ω := (0.05, 0.955).
In the two-bit environment, there exists another environment parameter a that controls the label noise level.
5

---

Published as a conference paper at ICLR 2023
Although the train-time bias parameters {0.1, 0.2) belong to Ωt test data is gener- ated afresh, different from training data. We see in FA that the superior- ity of IRM-GAME B 0.9 vanishes for smaller B. Consequently, for invari- ant prediction evaluated in other testing environments (e.g., 3 < 0.4), the per- formance ranking of the same methods becomes FRIVI= IRM-H---RRM--MMEE GAME. This mismatch of results suggests we measure the 'invariance' of IRM methods against diverse test environments. Oth- erwise, evaluation with single B could give a false sense of invariance. In Fig- ure[B,, present the box plots of pre- diction accuracies for IRM variants, over the diverse of testing environments (3 E ?2.. Evidently, IRMv1, the oldest (sub-optimal) IRM method, yields the least variance of invariant prediction accuracies and the best average prediction accuracy. compared to both IRM- TRME and FISHR. 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.
100
Incarinnt Accuracy (%)
Invariant Acouray
Toal Bax
40
20
&
Fair
0.0
0.2
0.4
0.6
11.8
1.0
50
Test Bay. Dias Parameter B
18Mvi
IRM-Gume
Fistrr
(A)
(B)
Figure 2: Performance comparison of TRM variants IRMv1, IRM-GAME. and FISHR COLORED-MNIED
(A) Evalua-
tion in different test-time environments (correspouding to dif- ferent 9). 3 values used by the two training env environments are 0.1, 0.2 respectively. The conventional evaluation is done with the test environment B = 0.9 (sec" 'a). (B) Box plots of pre- diction accuracies over diverse test environments correspond- ing to (0.05, 0.1,.... 0.95). IRMv1 achieves the best average accuracy (67. (67.13%), followed by FISHR (67.05%) and IRM-GAME (65.53%)
5 ADVANCING IRM-GAME VIA SANSENSUS-CONSTRRANNE
BLO
In this section, we revist and advance a special IRM variant, IR.M-GAm (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-GAME. We first introduce the setup of IRM-GAME following notations used in Section The most essential difference between IRM-GAME and the vanilla IRM framework is that the former assigns each environment with an individual classifier w( and then relies on the ensemble of these individual predictors, i.e., NV' Eeeg., (w(+) 0), for inference. IRM-GAME is in sharp contrast to [RM, where an environment-agnee the losses across all environments. Therefore, we raise the following question: Can IRM-GAME learn an invariant predictor?
prediction head w simultaneously optimizes
Inspired by the above question, we explicitly enforce invariance by imposing a consensus prediction w(N) and integrate it with IRM-GAME. Here, (e)) denotes the prediction head for the e-th environment. Based on the newly-introduced constraint, the ensemble prediction head X Ecer, (( can be interpreted as the average consensus
constraint C : ((w(1) w/2), .... w((h)) w((1)
over N environments: w" == N Σeee, w-(e) = arg min (w()) J. Eese,, [[w(e) w( I|2- With
the above consensus interpretation, we can then cast the invariant predictor-baked IRM-GAME as consensus-consstrained BLO problem, extended from (RRM):
minimize
Ececo ( (w)((0)) (0)
subject to ()) w (0) E arg min (()) (w (c))) -0), Ve Ex,
w(e)
(II) (0) Exer (0))
The above contains two lower-level problems: (I) per-environment risk minimization, and (II) pro- jection onto the consensus constraint ((w(e)) C). The incorporation of (II) is intended to ensure the use of invariant prediction head (0) in the upper-level optimization problem of (1).
Limitation of I and BLOC-IRM. In ( the introduced lem might compromise the optimality of the lower-level solution (0) to the per-environment (unconstrained) risk minimization problem (I), ie., violating the per-environment stationarity
lower-level prob-

---

Published as a conference paper at ICLR 2023
justifies this side effect. As stationarity is hardly attained at the consensus prediction when solving ( This is not surpris- ing since constrained optimization solution might not be stationary solution to minimizing the
W (w" (0) 0)1 Figure
the per-environment
(unconstrained) objective function. To alleviate this limitation, we improve ① by explicitly pro- moting the per-environment stationarity w/ (w" (0) 0)1B in its upper-level problem through optimization over 0. This leads to BLOC-IRM (BLO with Consensus IRM):
minimize Eeeru [w' (0) 0) W
0)1
(BLOC-IRM)
subject to Lower-level problems (I) and (II) in
where 7 > 0 is a regularization parameter like IRMv0 Assisted by the (upper-level) predic- tion stationarity regularization, the consensus prediction (II) indeed simultaneously minimizes the risks of all the environments, supported by the empirical evidence that the convergence of [[V (w" (0) 0)1 towards 0 along each environment's optimization path (see Figure A3).
can be effectively solved using 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 (1) is unconstrained and separable over environments, and the consensus operation (II) is linear. Based on these characteristics, the implicit
Further, elaborate how the BLOC-
gradient dw (0) can be directly computed as
dw (0)
1
Σ CE
dw (
subject to (())(0) E arg min ((5) (w (c) о 0)-
(2)
d0
N
do
(c)
Since the above lower-level problem is un- constrained, we can call the standard arg min differentiating method, such as implicit func- tion approach (Gould et al. 2016) or gradi- ent unrolling (Liu al.J (021) to compute (0) In our work, we adopt the gra- do
Backwardd
Algartthm
(e)
(b)
(a
Ixputi
Fonure Extrectrr
W"
dient unrolling method, which approximates w(())()) by a K-step gradient descent solu-
Classitios
(e) (0) and then leverages au- tion, noted by w tomatic differentiation (AD) to compute the derivative from w K
Predectoo Logiu
y
AP
Tang
yo!
you
(e) (0) to the variable 0. Fig- 3 shows the working pipeline of BLOC- IRM and its comparison to original IRM and IRM-GAME 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 the existing literature. We show the best performance is always achieved when the stationarity is penalized in the upper-level (see A3)
Figure 3: Schematic overview of BLOC-IRM over training environments (red and green), and compari- son to IRM and IRM-GAME.
EXPERIMENTS
In this section, we begin by introducing some key experiment setups (with details in Appendix C.I and then empirically show the effectiveness of our proposed IRM training and evaluation improve- ments over existing IRM methods across various datasets, models, and learning environments.
6.1 EXPERIMENT SETUPS
Datasets and models. Our experiments are conducted over 7 datasets as referenced and shown in Tables I 2 Among these datasets, COLORED- -NIST, COLORED-FMNIST. CIFAR-NNIST, and COLORED- OBJECT are similarly curated, mimicking the pipeline of COLORED- MNIST (Arjovsky et al. 2019), by introducing an environment bias parameter (e.g., 3 for COLORED -NNIS Sec- tion 4) to customize the level of spurious correlation (as shown in Table in different environments. In the CELEBA dataset, we choose the face attribute 'smiling' (vs. mnon-mmiinng') as the core fea- ture aimed for classification, and regard another face attribute 'hair color' ('blond' "dark') the
7

---

Published as a conference paper at ICLR 2023
source of spurious correlation imposed on the core feature. By controlling the level of spurious cor- relation, we then create different training/testing environments in CELEBA. 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 Guurajann & Lopez-Paz (2020) that for these datasets, ERM could even be better than IR Mv1. Yet, we will show that our (B..OO is a promising domain generalization method, which outperforms all the IRM baselines and ERM in practice. In addition, we follow Arjovsky et al. (2019) in adopting multi-layer perceptron (MLP) as and COLORED-FMNN complex datasets, we use the ResNet- 18 architecture (He et al. 2016).
the model for resolving COLORED-MNID
problems. In the other more
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 COLORED- -NEIST and COLLORDD-MMNE and 256 for other datasets. Sec-
tion 6.2 below, we also do thorough comparison of large-batch small-batch IRM training.
Evaluation setup. As proposed in Section 4 we use the multi-environment evaluation met- ric 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 mea- sured for IRM methods. The resulting perfor- mance is reported in the form a + b, with mean a and standard deviation computed across 10 independent trials.
Table Performance existing IRM methods in large and small- settings. GRAYSCALE refers to ERM on uncolored data, which yields the best predic- tion (supposing spurious correlation during train- ing). The IRM performance is evaluated by average accuracy ('Avg Acc') and accuracy gap ( ('Ace Gap') ) in the format mea A higher Avg Acc and lower Ace Gap is preferred. The theoretically optimal per-
formance 75% (Arjovsky el al. 2019).
Doniset
COLORED- MNIST
COLOREDD FMNNST
Matrics 5%)
Aug (t)
Acc ())
Avg (to
Acc ((p
GRAYSCALE
73.39
0.52
74.05an.1
0.136004
ERM
49.1921.50
90.72a2.06
49.7741.71
88.022249
6.2 EXPERIMENT RESULTS
Large Batch
IRM-GIMMT
IRMv1
IRMvO
65.39
20.53
8.75
4,69
3,4320
67.1920.22
65.91:00.9
66.4420
3,74:0000
3.3520
3.53:0115
REX
3.76:007
3.26
9IRM
+0.51
3.81
67.75::0.26
381:00111
Small-batch training improves all existing
SPARSHIRM
67.73 00.22
3.651008
89
3:12::11
IRM methods on COLORED-MNIST &
FISTIR
67.49:0.99
4.37
67.53:1921
4.49::016
FOL.ORTD-MMIIST. Recall from Section
3
Batch
Sinall Batch
IRMVE
IRMVO
68.3360:31
68:37:0028
2.04
1.320000
69.07.0027
1:45:000
L36:0000
IRM-GAME
67.73:0.21
1.6710.14 14
67.4950.52
L882:015
that all the existing IRM methods (Table adopt full-batch IRM training on COLORED-MNIST & COLORED-FMNIS' batch training problem. In Table 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 GrM-graascale (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.T the
REX
68-42:0.22 20.
1.65
68 28.66
1.29
RIRM 68.7140021
1.35
68 08.6440332
14440 15
SPARSEIRM
08.81:90.55
L.720005
68.29::022
1.2820
which raises the large-
PISUR
68:69:0019
2.13::008
68.7950017
1.77:0.10
IRM performance is measured by the average accuracy and the accuracy gap across 19 testing environments, parameterized by the environment bias parameter 3 E {0.05,... 0.95). 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-GAME accuracy to 8 in Figure 2 Third, we observe that IRMvO, 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, TableEjdemnnstrates the ef- fectiveness of our proposed BLOC-IRM approach versus ERM and existing IRM baselines across all the 7 datasets listed in Table Evidently, BLOC-IRA IRM yields a higher average accuracy com- pared to all the baselines, together with the smallest accuracy gap in most cases. Additionally, we observe that CELEBA, PACS and VLCS are much more challenging datasets for capturing invari- ance through IRM, as evidenced by the small performance gap between ERM and IRM methods. In
8

---

Published as a conference paper at ICLR 2023
Table 5: IRM performance comparison between BLOC--RM and other baselines. We use ResNet- 18 (He ct al. 2016) for all the datasets. The evaluation setup is consistent with T and the best performance per per-dataset is highlighted in bold. Wc present the results with the full dataset list in Table
Algorithm
COLORRN-ORIEC
CIPAR-NNIST
CELER
VICS
PACS
Metriss (9)
Avg Ace
Acc Cup
Arg Ace
Acc Cup
Arg Ace
Acc Cup
Arg Ace
Ace Gup
Avg Auc
Ace Cup
ERM
41.11:144
26.13
40.39-1.32
25.53
72.38-00.2
10.73.0.36
03.
12.39
69.95
14.32.0.75
1RMv1
64.42-0.21
4.18+0.29
61.49-0.29
7.17+0.33
02.49-0.36
10 10.15+0.27
62.72-0.25
12.74-0.27
68:93-00:5
14:99-0.51
IRMV
22.39.0.25
5.36100311
60.14-0118
8.83100.30
72.42-0.35
10.43.0.35
0.
12.99.0.36
68.72.10.29
15.29-0.71
IRM-GAME
62.88:0.34
5.5910:28
60.44-0.5
6.72:0.41
72. 12.18:0.44
12.32:00:41
6.
13.37
68.
68.12:30.22
15.77:0000
REx
63.37-0.35
5.42+-0.31
62.32-0.31
5.55+00.52 52
72. 02.34-0.36
10.31+0.22
63. 03.59
12.87-0.51
69.45-0.511
15.51-0.67
BIRM
65.11
3.312:022
62.99_0..35
5.2310.30
12.93.0.25
0.92:0.33
63 63.33
12. 12.13::0.23
69.31
15,76200.9
SPARSEIRM
64.9720.33
3 3.97:0022
62. 0.22
4.14::.311
72.42-00:3
9.79.0.44
62. 88.86...22
12.79.0.33
69.52.10:39
15.81:0.82
FISHE
64.77-0.25
4.41+-0.20
61.79-0.22
5.55+-0.21
9.42+-0.5
11:93-0049
14,52-0:43
BLOC IRM
4,8930030
73.35
8.79.00.21
11.5520.32
70.31:00221
14.7530008
particular, all the IRM methods, except FISHR and BLOC-IRM, could even be worse than ERM on PACS and VLCS. Here, we echo and extend the findings of [fueerr et al.] (202T] Section 4.3). However, we also show that BLOC-IRM
quite competitive IRM variant when applied re- alistic domain generalization datasets. We also highlight that the CELEBA experiment is newly constructed and performed in our work for invariance evaluation. Like PACS and VLCS. this ex- periment also shows that ERM is strong baseline, and among IRM-based methods, BLOC-IRM the best-performing, both in terms of accuracy and variance of invariant predictions.
IRM against model size and training environment variation. Furthermore, we investigate the effect of model size and train- ing environment diversity
Invariant Accuray
68
the IRM performance. The recent
works (Lin et al.] 2022 Zhou et al.] 2022b) have empirically shown that IRMv1 I may suffer significant performance loss
66
64
when trained over large-sized neural network models, and thus
1834-V
BDC-V
SIRM
developed BIRM and SPARSEIRM approaches as advancements of IRMv1. 1. Inspired by these works, Figure44presenn the sensi- tivity of invariant prediction to model size for different IRM meth-
62
IBM Gome
Fistr
BLOC-ERM
200
400 600 800 1000
Intermediate Layer Dimension
ods on COLORED-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 fol- lowed in the subsequent literature. As we can see, when d > 390, nearly all the studied IRM methods (including BLOC--RMM) suf- fer performance drop. Yet, as d 2 800, from the perspective of prediction accuracy and model resilience together, the top-3 best IRM methods with model size resilience are BIRM, SPAR- SEIRM, and BLOC-IRM, although we did not intentionally de- sign BLOC-IRM IIM to resist performance degradation against model size.
Figure 4: IRM performance on COLORED- layer dimension in MLP. The dot- ted line represents the default di- mension (d 390) used in the literature. The invariant prediction accuracy is presented via the dot line (mean). The results are based on 10 independent trials and we port the variance in Figure/A4]
against the
We also show more experiment results in the Appendix. In Table|A2] we study IRM with different numbers of training environment configuran configurations and observe the consistent improvement of BLOCC- IRM over other baselines. In Table A4 we show that the performance of invariant prediction de- grades, if additional covariate shifts (class, digit, and color imbalances on COLORED-MNIST) are imposed on the training environments following Krueger et al. (2021) Section 4.1) and also demon- strate that BLOC-IRM maintains the accuracy improvement over baselines with each variation. In Table we compare the performance of different methods in the failure cases of IRM pointed out by (Kamath et al. 202T) and show the consistent improvement brought by BLOC-IRM.
7 CONCLUSION
In this work, we investigate existing IRM methods and reveal long-standing but chronically over- looked 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 multi- environment evaluation. We reexamine the IRM-GAME method through the lens of consensus- constrained BLO, and develop a novel IRM variant, termed BLOC-IRM. We conducted extensive experiments 7 datasets and demonstrate that BLOC-IRM consistently improves all baselines.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENT
The work 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 CIF- 1910385. The computing resources used in this work were partially supported by the MIT-IBM Watson AI Lab and the Institute for Enabeee Research (ICER) at Michigan State University.
REPRODUCIBILITY STATEMENT
The authors have made an extensive effort to ensure the reproducibility algorithms and results pre- sented in the paper. First, the details of the experiment settings have been elaborated in Scction and Appendix(C.] 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 in- troduced in Section Second, eight IRM-oriented methods (including our proposed BLOC-IRM M) are studied in this work. The implementation details of all the baseline methods are clearly presented in Appendix C.2 including the syyper-parametes tuning, model configuration, and used code bases. For proposed BLOC-IRM,,
include all the implementation details in Section and Ap- pendix[B.3 including training pipeline in Figure and the pseudo-code Algorithm/A1 Third, all the results 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: /ttptttttuuuu.go http:/u..c.om//PPP.---G/G-O -ww.///PPPI----OO//111--------- TRM
10

---

Published as a conference paper at ICLR 2023
REFERENCES
Faruk Ahmed, Yoshua Bengio, Harm van Seijen, and Aaron Courville. Systematic gencralisation with group invariant predictions. International Conference on Learning Representations, 2020.
Kartik Ahuja, Karthikeyan Shanmugam, Kush Varshney, and Amit Dhurandhar. Invariant risk minimization games. In International Conference on Machine Learning, pp. 145-155 155. PMLR, 2020.
Kartik Ahuja, Wang, Amit Dhurandhar, Karthikeyan Shanmugam, and Kush Varshney. Empirical or invariant risk minimization? a sample complexity perspective. In International Conference on Learning Representations, 2021.
Martin Arjovsky, Léon Bottou, Ishaan Culrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv:1907.02892 2019.
Yogesh Balaji, Swami Sankaranarayanan, and Rama Chellappa. Metareg: Towards domain generalization using meta-reguularization. Advances in neural information processing systems, 31, 2018.
Sara Beery, Grant Van Horn, and Pietro Perona. Recognition in terra incognita. In Proceedings of the European conference computer vision (ECCV), pp. 456-473, 2018.
Francesco Cappio Borlino, Antonio D D'Innocente, and Tatiana Tommasi. Rethinking domain generalization baselines. In 2020 25th International Conference on Pattern Recognition (ICPR), pp. 9227-9233. IEEF, 2021.
Fabio M Carlucci, Antonio D' I'Innocente, Silvia Bucci, Barbara Caputo, and Tatiana Tommasi. Domain gener- alization by solving jigsaw puzzles. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 2229 2238, 2019.
Shiyu Chang, Yang Zhang, Mo Yu. Tommi Jaakkola. Invariant rationalization. International Conference on Machine Learning, pp. 1448-1458, PMLR, 2020.
Yongqiang Chen, Kaiwen Zhou, Yatao Bian, Binghui Xie, Kaili Ma, Yonggang Zhang, Han Yang. Bo Han, and James Cheng. Pareto invariant risk minimization. arXiv preprins uxXiv:2206.07766, 2022.
Yo Joong Choe, Jiyeon Ham, and Kyubyong Park. An empirical study of invariant risk minimization. arXiv preprint arXiv:2004.05007 2020.
Pim De Haan, Dinesh Jayaraman, and Sergey Levine. Causal confusion imitation learning. Advances in Neural Information Processing Systems, 32, 2019.
Alex DeGrave, Joseph D Janizek, and Su-In Lec. Ai for radiographic covid-19 19 detection selects shortcuts over signal. Nature Machine Intelligence, 3(7):610-6)) 619, 2021.
Qi Dou, Daniel Coelho de Castro, Konstantinos Kamnitsas, and Ben Glocker. Domain generalization via model-agnostic learning of semantic features. Advances in Neural Information Processing Systems, 32, 2019.
Yana Dranker, He He, and Yonatan Belinkov. Imm-when it works and when it doesn't: A test case of natural language inference. Advances in Neural Information Processing Systems, 34:18212-188 18224, 2021.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for effi- ciently improving generalization. arXiv preprint 4rX1v:2010.01412, 2020.
Yaroslav Ganin, Exgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Donaiinadversarial training of neural networks. The journal of machine learning research, 17(1):2096-230) 2016.
Robert Geirhos, Jöm-Henme Jacobsen, Claudio Michaclis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix Wichmann. Shortcut learning in deep neural networks. Nature Machine Intelligence, 2(11); 665-673, 2020.
Stephen Gould, Basura Fernando, Anoop Cherian, Peter Anderson, Rodrigo Santa Cruz, and Edison Guo. On differentiating parameterized argmin and argmax problems with application to bi-level optimization. arXiv preprint 1607.0006 2016.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch SGD: Training imagenet in 1 hour. arXiv preprint arXiv: 000../0066666000 02677, 2017.
11

---

Published as a conference paper at ICLR 2023
Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization. arXiv:2007.01334 2020.
arXiv preprint
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Neal Jean, Marshall Burke, Michael Xie, W Matthew Davis, David B Lobell, and Stefano Ermon. Combining satellite imagery and machine learning predict poverty. Science, 353(6301)77907944 2016.
Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Domain extrapolation via regret minimization. arXiv preprint arXiv:2006. 03908, 2020.
Pritish Kamath, Akilesh Tangella, Danica Sutherland, and Nathan Srebro. Does invariant risk minimization
capture invariance? In International Conference on Artificial Intelligence and Statistics, pp. 4069-4077
PMLR, 2021.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang- On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv:1009,0 04836, 2016.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint
arXiv: 1412.6980 2014.
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, et al. Wilds: A benchmark of in-the- wild distribution shifts. In International Conference on Machine Learning, pp. 5637-5664. PMLR, 2021.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. Communications of the ACM, 60(6):84 90, 2017.
David Krueger, Ethan Caballero, Joern-Hearik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. Out-of-distribution generalization via risk extrapolation (rex). In Interna- tional Conference Machine Learning, pp. 5815 5826. PMLR, 2021.
Da Li, Yongxin Yang. Yi-Zhe Song, and Timothy M Hospedales. Deeper, broader and artier domain general- ization. Proceedings of the IEFE international conference on computer vision, pp. 5542-5550, 2017.
Da Li, Jianshu Zhang, Yongxin Yang, Cong Liu, Yi-Zhe Song, and Timothy M Hospedales. Episodic training for domain generalization. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 1446-148 1455, 2019.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. Advances in neural information processing systems. 31, 2018.
Yong Lin, Qing Lian, and Tong Zhang. An empirical study invariant risk minimization on deep models. In ICML 2021 Workshop on Uncertainty and Robusiness in Deep Learning, pp. 7, 2021.
Yong Lin, Hanze Dong, Hao Wang, and Tong Zhang. Bayesian invariant risk minimization. In Proceedings of the IEEF/CVF Conference Computer Vision and Pattern Recognition, pp. 16021-16030, 2022.
Risheng Liu, Jiaxin Gao, Jin Zhang, Deyu Meng, and Zhouchen Lin. Investigating bi-level optimization for learning and vision from unified perspective: A survey and beyond. IEEE Transactions Pattern Analysis and Machine Intelligence, 2021.
Ziwci Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceed- ings of International Conference on Computer Vision (ICCV), December 2015.
Mingsheng Long, Yue Cao, Jianmin Wang, and Michael Jordan. Learning transferable features with deep adaptation networks. In International conference on machine learning, pp. 97-105. PMLR. 2015.
Hyconscob Nam, HyunJae Lee, Jongehan Park, Wonjun Yoon, and Donggeun Yoo. Reducing domain gap by reducing style bias. Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recogni- tion, pp. 8690-8699, 2021.
Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen. Causal inference by using invariant prediction: identi- fication and confidence intervals. Journal of the Royal Stotistical Society: Series B (Statistical Methodology), 1855):947-1012, 2016.
12

---

Published as a conference paper at ICLR 2023
Alexandre Rame, Corentin Dancette, and Matthieu Cord. Fishr: Invariant gradient variances for out-of- distribution generalization. In International Conference on Machine Learning, pp. 18347-18377. PMLR, 2022.
Elan Rosenfeld, Pradeep Ravikumar, and Andrej Risteski. The risks of invariant risk minimization. arXiv preprint arXiv:2010005761, 2020.
Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Perey Liang. Distributionally robust neural net- works for group shifts: On the importance of regularization for worst-case generalization. arXiv preprint MrXiv:/911.0877 2019.
Shiori Sagawa, Aditi Raghunathan, Pang Wei Koh, and Percy Liang. An investigation of why overparameteri- zation exacerbates spurious correlations. In International Conference Machine Learning, pp. 8346-8356. PMLR, 2020.
Harshay Shah, Kaustav Tamuly, Aditi Raghunathan, Prateek Jain, and Praneeth Netrapalli. The pitfalls of simplicity bias in neural networks. Advances in Newral Information Processing Systems, 33:9573-9585, 2020.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409. 1556, 2014.
Yi Sun, Xiaogang Wang, and Xiaoou Tang. Deep learning face representation from predicting 10,000 classes. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1891-1898, 2014.
Antonio Torralba and Alexei A Efros. Unbiased look datasct bias. In CVPR 2011, pp. 1521-1528. IEEE, 2011.
Eric Tzeng, Judy Hoffman, Ning Zhang, Kate Sacnko, and Trevor Darrell, Deep domain confusion: Maximiz- ing for domain invariance. arXiv preprint arXiv:1412.3474, 2014.
Jindong Wang. Cuiling Lan, Chang Liu, Yidong Ouyang, Tao Qin, Wang Lu, Yiqiang Chen, Wenjun Zeng, and Philip Yu. Generalizing to unseen domains: A survey on domain generalization. IEEE Trunsactions on Knowledge and Data Engineering, 2022.
Chuanlong Xie, Haotian Ye, Fei Chen, Yue Liu, Rui Sun, and Zhenguo Li. Risk variance penalization. arXiv preprint arXiv:2006.07555 2020.
Renzhe Xu, Xingxuan Zhang, Peng Cui, Bo Li, Zheyan Shen, and Jiazheng Xu. Regulatory instruments for fair personalized pricing. In Proceedings of the ACM Web Conference 2022, pp. 4-15, 2022.
Yilun Xu and Tommi Jaakkola. Learning representations that support robust transfer of predictors. arXiv
preprint arXiv:21 10.09940, 2021.
Yang You, Igor Gitman, and Boris Ginsburg Large batch training of convolutional networks. arXiv preprint arXiv:1708. (00.00888 2017a.
Yang You, Igor Gitman, and Boris Ginshurg. Scaling SGD batch size 32k for imagenet training. arXiv preprint /rxiv:/70e 03888, 2017b.
Yang You, 7hao Zhang, Cho-Jui Hsieh, James Demmel, and Kurt Keutzer. Imagenet training in minutes. In Proceedings of the 47th International Conference on Parallel Processing, pp. 1. ACM, 2018.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert 76 minutes. arXiv preprint arXiv: 190..00962.0.... 2019.
Dinghuai Zhang, Kartik Ahuja, Yilun Xu, Yisen Wang, and Aaron Courville. Can subnetwork structure be the key to out-of-distribution generalization? In International Conference Machine Learning, pp. 12356- 12367. PMLR, 2021.
Gaoyuan Zhang, Songtao Lu, Yihua Zhang, Xiangyi Chen, Pin- Yu Chen, Quanfu Fan, Lee Martie, Lior Horesh, Mingyi Hong, and Sijia Liu. Distributed adversarial training to robustify deep neural networks at scale. In Uncertainty in Artificial Intelligence, pp. 2353-2363. PMLR. 2022a.
Xingxuan Zhang, Linjun Zhou, Renzhe Xu, Peng Cui, Zheyan Shen, and Haoxin Liu. Nico++: Towards better benchmarking domain generalization. arXiv preprint arXiv:2204.0 08040. 2022b.
Kaiyang Zhou, Ziwei Liu, Yu Qiao, Tao Xiang, and Chen Change Loy. Domain generalization: A survey, IEEE Transactions on Pottern Analysis and Machine Intelligence, 2022a.
Xiao 7hou, Yong Lin, Weizhong Zhang, and Tong Zhang. Sparse invariant risk minimization. In International Conference on Machine Learning, pp. 27222-27244. PMI.R, 2022h.
13

---

Published as a conference paper at ICLR 2023
APPENDIX
A DATASET SELECTION
Compared to existing work, we expand the dataset types for evaluating the performance of different
IRM methods (see Table 2.. In addition the most commonly-used benchmark datasets COLORED- MNIST (Arjovsky et al. (019) and COLORED-FMNIST (Ahuja et al., 2020), we also consider the
datasets CIFSR-MNIST (Lin at. 2021 Shah al.) 2020) and COLORED-ORJJEC (Ahmed al.
2020 Zhang et al. 2021). which impose artificial spurious correlations, MNIST digit number and object color, into the original CIFAR-I0 and COCO Detection datasets, respectively. Furthermore,
we consider other three real-world datasets CELEBA (Liu et al. 2015). PACS (Li et al.) (017) and VLCS (Torralba Efros, 2011). without imposing artificial spurious correlations. Notably, CELEBA was first formalized and introduced to benchmark IRM performance. The recent work (Gulrajani & Lopez-Part 2020) showed that when carefully implemented, ERM could outperform IRMvI in PACS and VLCS. Thus, we regard them as challenging datasets to capture invariance.
For COLORED-OBIECE dataset, we strictly follow the setting adopted in (Lin et al. (022) to gen- erate the spurious features. For CIFAR-MNIST we use the class "bird" and "plane" in the dataset CIFAR as the invariant feature, while the digit "0" and "I" in MNIST as the spurious correlation.
CELEBA dataset is, for the first time, introduced to measure IRM performance. We select the attribute "Smmilig" as the invariant label and use the attribute "Hair Color" (blond and black hair) to create spurious correlation in each environment.
B IMPLEMENTATION DETAILS
B.1 DETAILS ON LARGE-BATCH OPTIMIZATION ENIIANCEMENTS
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 LSGDD-baked IRM variant IRM-LSGD.
IRM-LALR: Next, we adopt layerwise adaptive learning rate (LALR) in IRM training. Follow- ing You et al., 2019), we advance the learning rate scheduler by assigning each layer of neural network-basse 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:
0 0,, =(1
41,5
(A1)
where O denotes the i-th layer of the model parameters at iteration t, and u represents the first-order gradient of the corresponding layer-wise model parameters. We use 7(10r,,i1 = min maa( [ c cul) as the scaling factor of the adaptive learning rate dh We use ci = ||n,,
and C2 = 1 in our experiments.
Lastly, we leverage sharpness-aware minimization (SAM) to simultaneously mini- mize the IRM loss and the loss sharpness. The latter is achieved by explicitly penalizing the worst- case training loss of model weights when facing small weight perturbations. This yields wide minimum within a flat loss landscape. More specifically, the sharpness-aware loss can be formu- lated as:
IRM-SAM:
min (0),
where
(SAM (8)
max ((O + tc,,
(A2)
S
where the parameter perturbation e is subject to the perturbation constraint ||e||2 S p. When applied to IRM, replace the per-environment training loss with the SAM loss, and adopt the p 0.001.
B.2 BLOC-IRM IMPLEMENTATION
As described in Section the BLOC-IRA algorithm solves the IRM problem with two optimiza- tion levels. We use 1-step gradient descent to get the lower-level solution. We retain the gradient
14

---

Published as a conference paper at ICLR 2023
graph in PyTorch to enable auto differentiation. We assign each of the classification head {w(o) separate optimizer and use the same learning rate as the feature extractor 0, For COLORED-MNIST and WOLORED-FMIISSS we adopt learning rate of 2 10 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 penalty weight of 10 as IRMvI and IRMvO.
Algorithm A1 BLOC-IRM
1: Initialization: Training data {x(*) from N environments, Model feature extractor Oo. and N
model classification heads {w]"),
(e) learning rate }m,} series, penalty weight {re) series.
2: for Step t = 0, .....,. do
3: Lower-level: update classification head for each environment:
((ee) ++1 w/ de(t) (w ②
0)
Ve E Etr,
(A3)
dw
A o W.w.
4: Consensus projection: Ve E Eie, W wi+1 Eese,
W
r+1
(c)
5: Upper-level: update feature extractor with stationary penalty:
0.+1 = 0: 71 Σ d (cle)
0) v (w
0)113)
(A4)
d0
0= 0=0.w
eEEr
6: end for
C EXPERIMENTATION
C.1 ENVIRONMENT SETUP
As proposed in Section we use the multi-environment evaluation metric unless specified other- wise. 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 MNNS FOLNRED--MNISTT, COI.ORED-OBJECT, CIFAR- MNIST, and CELEBA dataset, we manually create 19 environments with uniformly sampled bias parameter 3 E {0.05, 0.1, 0.95), where the environment bias parameter 8 controls the spurious correlation (see Section 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 the training environments, while we use the test set of all four environ- ments to form our proposed multi-environment invariance evaluation system.
C.2 BASELINES
For each baseline method, we follow its official PyTorch repository except IRM-GAME and SPAR- SEIRM. We translate the TensorFlow-based original code base of IRM-GAM to PyTorch. As one of the latest IRM advancements, the official code of SPARSEIRM is not yet publicly available. Therefore, reproduce SPARSEIRM in PoTor..
In particular, for COLORED- -MNIST and COLORED-FMNIST we stick to the original hyper- parameters 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, and 500 epochs in total, as suggested by the original IRMvI and inherited by its variants. For the small-batch setting, adopt the same penalty weight 10% Further, found that the warm-up
the penalty weight of 10 190 warm-up epochs,
15

---

Published as a conference paper at ICLR 2023
phase could be shortened without sacrificing accuracy. Therefore, use 50 warm-up epochs and total 200 epochs for all the methods.
For other datasets, adopt the batch size of 128 and use ResNet- 18 the default model architec- ture. 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 100th and 150th epochs.
C.3 ADDITIONAL. EXPERIMENT RESULTS
The influence of batch size with all the baselines. We show in Figure the influence of training batch size on the performance different methods. We observe Figure[AT] as Figure that full batch setting does not achieve the best performance, and the use of mini-batch (stochastic gradient descent) indeed improves performance.
70
Invariant Accuracy (%)
69
68
67
66
IRM-I
BIRM
65
IRM-Game
Fisls
64
64
2k
Batch Size
Figure Al: The performance of all the baselines in this work trained with different batch sizes on COLORED- MNIST dataset. The full data batch-size 50k. The invariant accuracy corresponds to the average accuracy evaluated based on the diversified environments-ble evaluation metric.
0.4
3.4
0.4
3.4
0.2
0.2
1.0
0.0
0.0
0.2
0.2
-0.4
-0.4
0.4
0.4
(A) Large Batch
(B) Small Batch
Figure A2: The loss landscapes of invariant prediction models acquired by (A) large-batch IRMv1 training with 50k batch size and (B) small- batch training with Ik batch size. The 2D loss landscape visualization is realized using the tool in (Li al. 2018). The x and represent the lincar interpolation coefficients over two directional vectors originated from the converged local optima. Here the numbers on the contour denote the loss values over test data.
Loss landscapes of IRMv1 with different batch sizes. We plot the loss landscapes of the models trained with IRMv1 on COLORED-MNN -NNIT using large (full) and small batch in Figure Using small batch training, IRMv1 FFggAAB] converges
smooth neighborhood of local optima. This also corresponds to flatter loss landscape than the landscape of the large-batch training (Fig- A2e A)). The loss landscapes demonstrate consistent results other experiments discussed in Section
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 ( and [BLCC-IMM WELLORDD-MNIS' For (BLOC-IRM) we use the regularization term 108. which is aligned with the penalty coefficient used in IRMv1. As
see, without the stationarity regularization,
16

---

Published as a conference paper at ICLR 2023
the stationary loss remains 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.
10²
10'
10"
Stationary
10
(2),
10
(2) Env2
BLOC-IRM, Env1
10
BLOC--RM, Eav2
75
150 200
Epoch Number
Figure A3: The per-environment training trajectory the stationarity loss of [ and (BLC--IMM) on COLORED-MNIST. The training setting is the same as Figure The algorithmic details can be found in Appendix B
Performance of all the methods with full dataset list. We show in Table the results all the methods on the seven datasets we studied. To be more specific, in Table we append the results of and COLORED-FMNISS into Table55aas 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.
COLORED-MMNE
Table A1: IRM performance comparison between our proposed BLOC-IRM method and other baselines under the full list of datasets. We use MLP for COLORED -ENES and BELORRDDFFNIIS and ResNet 18 (ee et al. 2016 for the rest datasets. The evaluation setup is consistent with Table and the best performance per-evaluation metric and per-dataset is highlighted in bold.
Cincernn MINIST
Checesnn FRNNSTT
Donmann Derec
Coos MNIST
PACS
407721771
123001
BBAS11
06.3349 31
(8.664331
1.45.4 29
6442463)
61 33.404333
3114433
6372-000
55 23
14.944331
LRM Gros
07.5249.3
182249
03384404
02461131
6724644
1233-643
13574843
12432555
15.774400
1.3900
4337466
08.344335 55
17
08.344330
1.3899 55
04574636
1974035
1240-110
630-406
12 12784825
15
BLOC-IRM
09:3320221
8899200
4.0000
62891132
7339
8391532
18551833
38.301821
Experiment on different model sizes. We show in Figure the influence of the increasing model size on the performance of different baselines considered this work. Compared Fig- 4 we report additional standard deviation of the 10 independent trials in A4
70
Invariant Accuracy (%)
68
66
64
IRM-VI
BIRM
IRM- VO
SIRM
62
IRM-Game
Fisha
REx
BLOC-IRM
60
200
400 600 800 1000
Intermediate Layer Dimension
Figure A4: IRM performance on COLORRD--MIIS against the dimension of the intermediate layer in MI.P The dotted line represents the default dimension (d = 390) used in the literature. The invariant prediction accuracy is presented via the dot line (mean) and shaded area (standard deviation) over 10 random trials.
17

---

Published as a conference paper at ICLR 2023
Experiment with different training environments. In Table we show the performance of all the methods in more complex training environments, such as more training environments and more skewed environment bias parameter B. As see, ILOC-IRM outperforms other baselines.
Table A2: Performance under different training enviromments in WOLORED-MNISH
Environment Metries (%)
Per E {0.1, 0.15)
Per E 0.15, 0.2)
Avg Acc
Acc Gap
Avg Ace
Ace Gap
OPTIMUM
75.00
0.00
75.00
0.00
GRAYSCALE
73.82400.1
0.37++0.05
73.97=0011
0,29+0.08
ERM
49.21440.79
91.88+3.31
49.03+-093
92.17443.04
IRMv1
67.36+40331
2.77:00.15
67.11=0.34
2.42+0.
IRMVO
67.01
2.85+00118
66.71=0:12
2.3620.19
IRM-GAME
66.39440.22
4.47+-0.11
65.93+0053
4.25++0.84
REX
66.82440444
2.59++0.11
67.14=0.38
2.16+0.118 13
BIRM
67.35++0.39
2.65:+0110
68.05=0.43
1.9920.07
SPARSEIRM
67.1240053
2.334:0.18
67.72=0041
2.11
FISHR
63.22240.4
2.444-0115
67.32-0039
2.59+0.15
BLO-IRM
68.72+0.41
2.19::0:15
68.89=0331
2.39:0.09
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 regulariza- tion, including the ones used in IRMv1 (BLOC-IRM-VI) REx (BLOC-IRM-REX, and FISHR 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 achieve an optimal solution.
Table A3: The performance of BLOC-IRI with different regularization terms. Three datasets studied and the latest baseline SPARSHIRM listed reference for comparison. The best performance per-evaluation metric and per-dataset is highlighted bold.
Dataset
COLORED-MNISS
COLOORD--OBECT
CIFAR-MNIST
Metries
Avg Ace
Acc Gap
Avg Ace
Acc Gap
Avg Acc
Acc Gap
SPARSEIRM
68.81
1.72:10.05
64.97:10.30
3.97-0.25
62.87:1.299
4.1410.31
BLOC-IRM
69.4710.21
1.0410.07
65.9710.33
4.10=0.36
63.6910.92
4.8910.36
BLOC-IRM-VI
67. 67.14:2.21
4.33=20.83
63.38:20.25 29
6.31=0.511
61.13:0.51
6.71
BLOC-IRM- REX
62.71110.1
8.74
60.3111033
7.62=0.66
59.3910.55
7.8910.45
BLOC-IRM-FISIR
63.25=0016
7.12:t0.39
61.17+0.31
6.98.22045
60.86:t0.1
6.63++0.30
Performance comparison of different methods with additional covariate shifts. Besides the sensitivity check on model size, Table/A4 BableeA4exammines examines the resilience of IRM to variations in the train- ing 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 COLORED-MNEL -NIIST by introduc- ing class, digit, and color imbalances, following Krueger et al. (202T] Section 4. Compared with Table IRM suffers greater performance loss in Table [44 in the presence of training environ- ment variations. However, the proposed BL.OC-IRM maintains the accuracy improvement over baselines with each variation. In Table A2 environments and observe the consistent improvement of BLOC-IRM over other baselines.
we also study IRM with different numbers of training
Exploration on the failure cases of previous IRM methods. Some papers either theoreti- cally (Rosenfeld et al. 2020) or empirically (Kamath et al. 2021) pointed out that the origi- nal IRMvi method could fail in certain cireumstances, due to the fact that the regularization term used in IRMv1 heavily relies on the "linear predictor" assumption. Regarding this is- sue, 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 IRMv0 in the upper-level objec- tive, not IRMv1. To justify our argument, we repeat the experiments in (Kamath et al.) (021), which points out a specific scenario using the COLORED-MNIST dataset where IRMv fails.
18

---

Published as a conference paper at ICLR 2023
Table A4: IRM performance on COL.ORED-MNIST and COLLORED-MMNIST with training environment variations in terms of class, digit and color imhalances. The best IRM performance per-evaluation metric and per-variation source is highlighted in bold.
Vartalion
Cleo. Inhtionce
Digi Incullnee
Oolor Inktalnne
Clas Imbelanse
Dapi Inbakenes
Caker Imbelsre:
Metrics (9)
Ang. Acc
Acc Ekg
Ang Not
Acc Cap
Arg Aoe
Gap
Arg
Apc Cup
Avg Aoe
Gip
Avg Acc
Acc Gap
GRAVSCAER
11:21:011
71.20101
1.80.1111
2.29:10111
09-9211815
71:11-0111
BRM
43.72etm
92.7081445
90/86-1.00
41.728106
93.3752.16 16
42,39+2219
02.2542:72
45.29:0227
91,31=2.27
IRM+
Wa
4.4800:25 25
66 81889-00300 06
110c 41
061235055
3.31-0.25
02.4940.55 50
4.91-0.45 0:
61 9660955
K.5000.0
61:99:0:11
3:79400:5
IRMvO
65.0130225
65.13em122 4.29
3.87:3326
06.7231.25
3.01
62.7825148 66
5.33
61/6281129
5,2921141
6493c6622
3.28210331
IRA--GAME
50:31 10.m
0:88-C-35
63 10:00-003
8.77-0044 01
61.521045
1:78-0055
A0-751055 64
6:24-0.5 43
60 00:70-0015
54710.82
64.13-0.11
5.7%-0311
RIX
66-25-0840 k3
10:31
66-911442
1.1310.3
64.89-18.5
61:911825
4.751008
6.87-..1
4.10::0:5
BIRM
4.1mm5.:
58
06.7240244
3.47
3,80=:0.29
62,0710355
5.551152
Fismn
395
05.87:9.40
3.72=9-41
03.4849
5,90e0.35
62.4750266
5.7240.8
4,46:031
BLOC-BR
96.329197
3.10:533
08:413123
3:30300155
07.2511031
3.7239127
65.9910331
65,1321031
5
66.79:0311
3/721106
More specifically, the models are trained the training environ- ments (a, = (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, IRMv1 1 is clearly worse than ERM as achieves much lower average accuracy and higher accuracy gap. However, BL.OC-IMM 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, fect 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 IRMv1, 1, but to fix the issue of IRM-Game that resorts to a predictor ensemble to make the invariant predic- tion, which deviates from the spirit of acquiring invariant predictors in the original IRM paradigm.
Table A5: Performance parisons on COLORED-MNED among ERM, IRMv1, and BLOCC-
IRM in the scenarios where IRM-
variants failed following (Kamath)
et al. 202T).
Method
Avg. Acc.
Acc. Gap
IRMVI BLOC-IRM
76.89
27.68
84.22
also acknowledge that BLOC-IRM is not per-
A similar curve to Figure 1 on COLORED-FMNIST., We show the results for COLORED- FMNIST similar to Figure in Figure/ASJad the conclusion does not change much. As mentioned before, the latge-batch training setup was typically used for IRM training over the COLORED- MNIST and COLORED-FMNISS datasets.
70
Invariant Accuracy (%)
69
68
67
IRM-VI
66
IRM-VO
65
REx
64
256
1k
2k
4k
50k
Batch Size
Figure A5: The performance of three IRM methods (IRMv1, IRMvO, and REx) vs. batch size under COL.RRED-MNNIST.
19


Reviews:
[{"review_id": "o0GHxV_jedN", "reviewer": "Reviewer_3u9X", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors first conduct an empirical study of IRM and the IRM Games method, finding that smaller batch sizes can lead to performance gains, and highlighting the importance of evaluating on multiple test environments. They then propose a new variant of IRM where an invariant prediction head is computed as the average of prediction heads from each environment. The authors evaluate their method on several datasets, finding that it outperforms the baselines.", "summary_of_the_review": "The paper presents a couple of interesting empirical results, as well as proposes a method with solid empirical performance. I am leaning towards accept, pending the authors' rebuttal.", "strength_and_weaknesses": "Strengths:\n- The paper is generally clear and well-written.\n- The proposed method outperforms the baselines. \n\nWeaknesses:\n1. The paper would be greatly improved with some theoretical justifications of the method. For example, in the original IRM Games paper, they showed that the set of predictors obtained with their method is the same as the one obtained via IRM. Is this also true for BLOC-IRM? In addition, Is the proposed method still susceptible to the failure modes of IRM explored in (Rosenfeld et al., 2020, Kamath et al., 2021)? \n\n2. The authors should apply their method to some real-world shifts from the WILDS benchmark (Koh et al., 2021).\n\n3. Regarding Section 4, though I agree that methods should be evaluated on a diverse set of test environments, I don't necessarily agree with just showing the average performance, as averaging performance over a large grid of test environments can hide performance issues in each one. For example, if most of the test environments are similar to the training environments, then this defeats the purpose of domain generalization. I think perhaps it makes more sense to only consider test environments that are sufficiently different from the training. The authors should also show the worst-environment performance in all the evaluations.\n\n4. The authors should explain why, in Figure 1, the performance with very small batch size (64) is greatly diminished. Is this an issue of compatibility with the other hyperparameters (e.g. learning rate, annealing steps), or is there another justification for this?\n\n5. The authors should show a similar curve to Figure 1 for all of the other datasets (perhaps in the appendix), as this is an interesting and important empirical result.\n\n6. The authors should clarify their model selection method, as this has been found to be extremely important factor in the performance of domain generalization methods (Gulrajani & Lopez-Paz, 2020). Presumably, the authors used the test domain to select their hyperparameters.\n\n7. The proposed method (Section 5) seems disconnected from the previous results, as the method formulation is not motivated by any of empirical findings there.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "kfvBbO0_h7w", "reviewer": "Reviewer_bEZc", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Invariant Risk Minimization (IRM) has recently attracted great attention for its invariance and causality. However, the performance of IRM is limited, and it shows less performance compared to the ERM. This paper raises practical solutions to solve the issues. This paper provides three solutions, batch size, validation set, and new IRM variants. The paper is well-written, and it has a persuasive discussion and experimental results.", "summary_of_the_review": "I read the paper three times carefully, and the paper is well written and novel.", "strength_and_weaknesses": "Strength\n1. This paper tackles the critical problems of IRM, and practical solutions to solve the issues. \n2. The solutions are simple and effective.\n3. The experiment sections are extensive, and they include diverse baselines.\n\nWeakness or Questions\n1. This paper provides the Block-IRM, the extension of the IRM game. Is it still possible to analyze the Block-IRM theoretically?\n2. This paper extends the IRM as IRM-SAM. I wonder about the author's opinion about the relationship between IRM and sharpness.\n3. This paper provides a diverse dataset, but the dataset is limited to the toy or the small-scale dataset. Validation on the large-scale dataset such as DomainNet will be helpful.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "CtNIGGfCfEY", "reviewer": "Reviewer_U4vc", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes improvements to the invariant risk minimization (IRM) framework that ought to make the IRM approach more successful in practice. Specifically, the paper proposes to train IRM classifiers using smaller-batch data, conducts a more comprehensive assessment of current methods to show that some of the previously claimed benefits regarding invariance to environments might not be the case, and lastly introduces a modification of the IRM-Game algorithm of Ahuja et al. Overall, the comprehensive experiments suggests that this approach is beneficial.", "summary_of_the_review": "This paper proposes 3 different modifications and checks to the current IRM paradigm. Overall, the paper provides comprehensive empirical justification for its claims; however, it is unclear how these modifications fit together.", "strength_and_weaknesses": "### Strengths\n**Comprehensive empirical demonstration**: for the three changes that this paper proposes to IRM, the authors perform comprehensive assessment of these changes and present empirical evidence that demonstrates improvement in each case. I'll speak to the scale of the improvement in the next section. However, we observe from Table 4 that across the board, small-batch training does indeed improve performance of the IRM variants. Table 5 and 6 show that the BLOC-IRM proposal indeed provides additional gains as well. All the values reported also have error bars as well, so the results should likely hold up to scrutiny.\n\n### Weaknesses\n- **BLOC-IRM Formulation**: First the BLO acronym is not really introduced. Second, it is unclear to me why this formulation addresses the challenge discussed in the paper by Rosenfield et. al. on the limitations of IRM. Further, it seems that this formulation would actually not side-step this issue since you need as many classifiers as you have environments. \n- **Small batch training**: While small-batch training does give performance improvements, it is unclear why it does so. In addition, these improvements seem quite marginal but consistent. Any thoughts as to why?\n- **Unifying theme**: It is unclear how the three portions of this paper fit together. \n- **Why IRM-Game**: Why did the authors choose the IRM-Game formulation? It seems like one of out several IRM variants\n\n\n___\nPost Rebuttal Notes\nThe response has addressed most of my concerns, so I am raising my score.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster