Environment Diversification with Multi-head Neural Network for Invariant Learning
Authors: Bo-Wei Huang
Keng-Te Liao
Chang-Sheng Kao
Shou-De Lin
Venue: NeurIPS 2022 Accept (2022)
Keywords: Invariant Learning
Out-of-Distribution Generalization
Deep Learning

Abstract: Neural networks are often trained with empirical risk minimization; however, it has been shown that a shift between training and testing distributions can cause unpredictable performance degradation. On this issue, a research direction, invariant learning, has been proposed to extract causal features insensitive to the distributional changes. This work proposes EDNIL, an invariant learning framework containing a multi-head neural network to absorb data biases. We show that this framework does not require prior knowledge about environments or strong assumptions about the pre-trained model. We also reveal that the proposed algorithm has theoretical connections to recent studies discussing properties of variant and invariant features. Finally, we demonstrate that models trained with EDNIL are empirically more robust against distributional shifts. 

Environment Diversification with Multi-head Neural Network for Invariant Learning
Bo-Wei Huang
Keng-Te Liao
Chang-Sheng Kao Shou-De Lin
Department of Computer Science and Information Engineering
National Taiwan University, Taiwan
,00992000,,, b b07902046 sddinn)ccsieeeeeeeeeeeed nnu edu. tw
Abstract
Neural networks are often trained with empirical risk minimization; however, it has been shown that shift between training and testing distributions can cause unpredictable performance degradation.
this issue, research direction, in- variant learning, has been proposed to extract causal features insensitive the distributional changes. This work proposes EDNIL, an invariant learning frame- work containing multi-head neural network to absorb data biases. We show that this framework does not require prior knowledge about environments or strong as- sumptions about the pre-trained model. We also reveal that the proposed algorithm has theoretical connections to recent studies discussing properties of variant and invariant features. Finally, demonstrate that models trained with EDNIL empirically more robust against distributional shifts.
Introduction
Ensuring model performance on unseen data is a common yet challenging task in machine learning. A widely adopted solution would be empirical risk minimization (ERM), where training and testing data are assumed to be independent and identically distributed. However, data real-world applications with undesired biases, causing shift between training and testing distributions. It has been known that the distributional shifts can severely harm ERM model performance and even cause the trained model to be worse than random predictions [0]. In this work, foeus Invariant Learning, which aims at learning causal features expected to be robust against distributional changes. Invariant Risk Minimization (IRM) [ has been proposed Specifically, IRM is based on an assumption that training data are collected from multiple sources or environments having distinct data distributions. The learning objective is then designed standard ERM loss function with penalty term constraining the trained model (e.g. classifier) to be optimal all the environments.
popular solution for invariant learning.
IRM has shown to be effective; however, we note that IRM and many invariant learning methods rely on strict assumptions which limit the practical impacts. The limitations summarized follows.
Prior Knowledge of Environments IRM assumes training data are collected from environments, and the environment labels (i.e. which data instance belongs which environment) are given. However, the environment labels are often unavailable. Moreover, the definition environments can be implicit, making human labeling more difficult and expensive. To find environments without supervision, Creager et al. 6 propose EIIL, a min-max optimization framework training the model of interest and inferring the environment labels. Another work, HRM 22 (or the extension, KerHRM 23 ). parameterizes the environments and proposes clustering based approaches to estimate the parameters. A recent method, ZIN 20), also learns to label data. However, it relies carefully chosen features satisfying a series of theoretical constraints, and thus human efforts are still required.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Delicate Initialization EIIL is able to infer environments but requires an ERM model for initializa- tion. Crucially, the ERM model should heavily depend on spurious correlations. Creager al. [ reveal that, for example, slightly underfitted ERM models may encode more spurious relationships in some cases. However, as the distributional shifts are assumed to be unknown in the training stage, appropriate initialization might be difficult to guarantee.
Efficiency Issue HRM and KerHRM, though do not possess the above two limitations, suffer from the efficiency issue. Specifically, HRM is assumed to be trained with low dimensional data. As for KerHRM, although extends HRM to avoid the issue of dimensions by adopting kernel methods, the computational costs of the proposed method can be very high if the data or model size is large.
This work proposes novel framework, Environment Diversification with multi-head neural Network for Invariant Learning (EDNIL). EDNIL able to infer environment labels without supervision and achieve joint optimization of environment inference and invariant learning models. The underlying multi-head neural network explicitly diversifies the inferred environments, which is consistent with revealing the benefits of diverse environments. Notably, the proposed neural network is functionally similar to a multi-class classifier and can be optimized efficiently. The advantages of EDNIL summarized
recent studies S
20 [
We implement this framework using various pre-trained models such as Resnet 14 and DistilBert 30I, and evaluate it with diverse data types and varied biases. The results show that EDNIL can constantly outperform the existing state-of-the methods.
The learning algorithm of EDNIL has theoretical connections recent studies discussing conditions of ideal environments.
20
EDNIL does not have the three limitations discussed above. The comparisons between EDNIL and other methods are shown in Table Il
Table I: A summary of the advantages of invariant learning methods.
Unsupervised Insensitive Initialization
Efficiency
IRM
ZIN
20
EIIL
HRM
22
X
KerHRM
X
EDNIL (Ours)
2 Preliminaries and Related Works
The goal of EDNIL is to tackle out-of-distribution problems with invariant learning in the absence of manual environment labels. In Section 2.1 background knowledge about out-of-distribution generalization and invariant learning are introduced. In Section 2.2 discuss recent studies investigating ideal environments. Scction 2.3] we introduce the existing unsupervised methods inferring environments.
2.1 Out-of-distribution Generalization and Invariant Learning
Following
22 we consider a dataset D = {D*) (eesspppp(E)) with different sources D
(fff, Hi collected under multiple training environments E supp(Eu). Random variable Er indicates the training environment labels. For simplicity, Xe, y and pe denote data, target label and distribution in environment respectively.
With Eal containing all possible environments such that supp(Ean) supp(En), the goal of out- of-distribution generalization is to learn a predictor f(-) X' + as Equation [] where (f) =
[X Y (xxxy=(((((')) E [(((),"") Y measures the risk under environment with any loss
Exe,y=
A method is unsupervised if it does require extra human efforts obtain environments.

---

function 1(..)). ). In general, for E supp(Eu) and e' supp(Ean) \ supp(EE)) p (X,Y) is rather different from pa (=X.)Y)
= arg min
(f)
eEsupp(ea)
have attempted to tackle the generalization problems by discovering invariant relationships across all environments. commonly proposed assumption the existence of invariant features X. and variant features X,. Specifically, raw features X are assumed be composed of X, and X,, X h(x.. Xx,) where h(-) is transformation function. Invariant features Xg are assumed to be equally informative for predicting targets Y across environments supp(Ean). On the contrary, the distribution P (Y(()x,) can arbitrarily vary across e. As result, predictors depending on X, can have unpredictable performance in unseen environments. In particular, the correlations between X, and Y are known as spurious and unreliable.
Recently, several studies 11 26 28
To extract X.. IRM [1] assumes there is an encoder & for obtaining representations 2(X) ~ Xe- The encoder is trained with a regularization term enforcing simultaneous optimality of the predictor w & in training environments, where dummy classifier w = 1.0 is fixed multiplier for the encoder outputs:
Σ
Rd + \|A 1. (w 4)
(2)
eEsupp(E.) (E.)
As w is a dummy layer, the encoder also regarded as a predictor.
2.2 Ideal Environments
As Eu is unavailable sub-optimal in most applications, learning find appropriate environments (denoted by &lam) is attractive. However, the challenge is the lack knowledge of valid environments. Recently, Lin et al. [20 have proposed Equation B and 4 as the conditions of ideal environments, where is conditional entropy. To satisfy the conditions, Lin et al. [20 propose leveraging auxiliary information for model training. However, the method still requires extra human efforts collect and verify the additional information.
H(Y = H (Y (Xo Eteam)
(3)
XY) HYY,X, Eleum))
>
(4)
Particularly, Equation 4 can be implied by empirical studies [5] where diversity of environments is recognized as the key obtaining effective IRM models. To more precise, large discrepancy of spurious correlations, clear indication of distributional shifts, IRM can easily identify and eliminate variant features. Beyond
Pe (Y ). between environments is favored. As the environments give
IRM, HRM [ and KerHRM 23 can also be viewed optimizing diversity via clustering- methods specifically.
2.3 Unsupervised Environment Inference
Here provide detailed introduction of the existing environment inference methods that do not require extra human efforts. The general idea is to integrate environment inference with invariant learning algorithms who require provided environments.
EIIL proposes formulating invariant learning as a min-max optimization problem. Specifically, EIII. is composed of two objectives, Environment Inference (FT) and Invariant Learning (TT.), where El optimized by maximizing the training penalty via labeling the data, and LL is optimized by minimizing the training loss given the data labeled by El. The two-stage framework bypasses the difficulty of defining environments; however, the training result heavily relies on the initialization of the E1 optimization. Specifically, the initialization demands strongly biased ERM reference model; otherwise, EIIL can have a significantly weaker performance.
Another method, HRM 22
proposes a clustering -bsed method for learning plausible environments. HRM assumes that spurious correlations in each environment can be modeled by a parameterized function and the dataset is generated by the mixture of the functions. The parameters are then estimated by employing EM algorithm. Additionally, HRM equips a joint learning framework which alternatively learns invariant predictors and improves quality of clustering results. However, known
3

---

Dataset
(F-- ronoccatle)
Initializze with ERM
Minimize
Minimiiz Eyz
MEI
Eiearn
MIL
Diversified Environments
Alternating Update
Minimize
Invariant Prodictor
Mg1
d
(a)
(h)
Figure 1: Concepts EDNIL. (a) The joint learning framework of EDNIL. The environment inference model ME, containing variant encoder V and environmental functions f", is trained with LEi = LED + 3L + Lpp The invariant learning model Mn. containing an invariant predictor 4, is trained with Lu. (b) The multi-head network structure of the environment inference model Ma..
issue of HRM is assumption that the data are represented by raw features. Data such as images and texts requiring non-linear neural networks to obtain representations are beyond the capability.
To extend HRM to a broader class of applications and improve the model performance, Liu et al. propose KerHRM. The main idea is adopt the Neural Tangent Kernel [I5] method which transforms non-linear neural network training into a linear regression problem on the proposed Neural Tangent Features space. result, KerHRM elegantly resolves the shortcomings of HRM and is shown to be more effective. However, the proposed method and its implementations bring additional computational costs depending on data and model capacity. For applications favoring large datasets and pre-trained models, such as Resnet 14 and BERT 7 KerHRM may not be affordable option at the present stage.
23
Methodology
In this section, we propose general framework learn invariant model without manual environment labels. As shown in Figure Ta our proposed method consists of two models, Mg1 and MI.. Given the pooled data Y), Mer infers environments &xam satisfying Condition 3 and and Mj, is an invariant model trained with the inferred environments. Our framework is jointly optimized with alternating updates. The learned Mj, can provide information of invariant features ME, so that Condition 3 and can be fulfilled simultaneously. Note that Ma only serves at train time to provide environments for invariant learning. At test time, only Mn. needed perform invariant predictions.
3.1 The Environment Inference Model
The target of environment inference is to partition data with environment labels Exem satisfying Condition and this regard, propose graphical model, which sufficient condition for Condition and (the proof is in Appendix AA) as our foundation of inference model and learning objectives. The graph is shown in Figure 2 where the data generation process follows the proposed example in
X₂
Y
EjearT
Figure 2: The graphical model.
The inference model, ME. a neural network learning to realize the underlying graphical model. Following the idea of parameterizing environments from HRM [22] and KerHRM [24], we assume the distinct mapping relation between X and Y in environment e can be modeled by function (XX),, where V/ (XX) is learned representations expected to encode variant features X, and is an environmental function responsible for predicting Y. Instead of employing clusters, we propose building a multi-head neural network as shown in Figure lb] particularly, a single-head network ME with shared parameters corresponds cluster center in HRM or KerHRM. The training procedure of Mg can be divided into inference stage and learning stage.

---

3.1.1 Inference Stage of ME1
The goal is to infer an environment label for each training data. As in the graphical model, Etearn associated with variant relationships. Inspired by multi-class classification problem, we propose Equation 5 where the probability P(e| X, X) is estimated via softmax of negative (()), Y) divided by a constant temperature т. The function is expected to be the commonly used cross entropy mean squared error that measures the discrepancy between Y and (X))) for each environment e. Intuitively, each data prefers the environment whose model has better prediction.
exp
Y)/T)
P(E = e X,Y) = Σ
(5)
exp (-1(f (V()),
Esupp(Eum))
3.1.2 Learning Stage of Mg
The goal is to update the neural network to improve the quality of inference. Based on the structure of the graphical model, three losses are designed for minimization, i.c. Environment Diversification Loss (LED), Label Independency Loss (Lu) and Invariance Preserving Loss (Lyp). In particular, LED and Lp correspond the concepts of Condition 4 anddd and respectivel...
Environment Diversification Loss (LED) We consider maximizing H Eterm)) to satisfy Condition and capture more diverse variant relationships. Given the estimated P( XX Y), LED selects the most probable environment and its corresponding network for optimization:
LED Σ w; max[log P(e I2, y/)]
(6)
For each data (x,, y:), although only one environment e; = arg max, P(e Tz Ui) is selected for the minimization, the softmax simultaneously propagates gradient to maximize [[fr² (Vxz)) y1 for é' # eg. The network learns to maximize the dependency between Eteam and Y given variant representations. In terms of spurious correlations, the distinction between environments is expected to become clearer accordingly. In practice, we utilize scaling weight wi inversely proportional the size of eq. The importance of smaller environments will be thus enhanced within the summation.
Label Independency Loss (Lμ) With d-separation P5 Exeam is independent of Y in the graphical model. Hence, Lu constraints their dependency measured by the mutual information I(Y; Sream). Empirically, Lu prevents trivial solution that environments determined purely by target labels Y regardless of input features X, which is undesirable for invariant learning. To minimize I(Y; Eteam). it can be verified that it is equivalent minimizing Equation 7 given P(Y) is known.
Lu P
P(yle) log P(yle))
g
Invariance Preserving Loss (Li)) For Condition 3 as Mr. learns some invariant relationships after several training steps (Section 3.2 Lip can be considered to exclude invariant features from the diversification. Specifically. designed as the contrary of LED. LIP limits the variance of expected loss from invariant predictor & (in MiL) across environments (Equation (8. Similar idea can be found in [77. However, instead of training invariant model given known environments, we freeze the invariant predictor and regularize the adjustment of Eteam (i.e. the updates of MEa) here.
LIP Var
(X)(XX) Y"II
(8)
In summary, the training loss of MET can summarized as Environment Inference Loss (LE)). The regularization strengths of LLI and Lip can be controlled by hyper-parameters 3 and respectively:
L LED + 3T + TTP
In addition, before minimizing LE. we treitrain our V and one arbitrary fr with ERM. In general, it empirically facilitates better feature extraction. Unlike EIIL 66 taking ERM as a reference model heavily relying on variant features, EDNIL performs more consistently under various choices of FRM. Namely, the initialization of EDNIL can be more arbitrary than that of EIII.. We verify the argument in Section 4
5

---

3.2 The Invariant Learning Model
To identify invariance across environments, IRM is selected as our base algorithm optimizing the invariant predictor & in our model Mi.. As for the required environment partitions during training, we assign environment label E supp(Eieam) with largest P(e|z,, Vi), inferred by MF (Section 3.1.I). to each data (#a, Vs). However, it is inevitable that there exist some noises in automatically inferred environments, especially in the beginning of joint optimization. To reduce the impact of immature environments on invariant learning, we calculate the confidence score Ce for each environment e E supp(Exam), i.c. ES [P(eeeX" Y=)]. Our training objective is modified to minimize Invariant Learning Loss (Ln,) that considers the weighted average of environmental losses:
Ln =
Σ
We ( + A/A 1f o)
(10)
@Esupp((xum)
We E Eeepppp(Eum)) Ce'
(11)
Experiments
We empirically validate the proposed method biased datasets, Adult-Confounded CMNIST, Waterbirds and SNLI. The generation of spurious correlations mainly follow the protocols proposed
by M6 9 29 In Section 4.1 Colult-Confoudee and CMNIST tested with Multilayer Perceptron two more complex datasets, Waterbirds and SNLI, are taken for evaluating the
(MLP). In Section integration of transfer learning. Deep pre-trained models will be fine-tuned to discover variant and invariant representations.
The following four methods are selected as competitors: Empirical Risk Minimization (ERM), Environment Inference for Invariant Learning (EIIL (6),, Kernelized Heterogeneous Risk Minimiza- tion (KerHRM [X3] and Invariant Risk Minimization (IRM (1] Equation 2 EIIL and KerHRM are invariant learning methods with unsupervised environment inference, which share the same settings EDNIL. HRM [22] is replaced by KerHRM for non-linearity. For IRM who requires environment partitions, we re-label Eoracle on each given biased training set, which diversifies the spurious relationships elicit upper-bound performance of IRM.
tuning, we split 10% of training data to construct an in-distribution validation set. In each dataset, several testing environments with different distributions are listed to evaluate the robustness of each method, and we mainly take worst-case performance for assessment. As all tasks in this section are classification problems, accuracy is adopted the evaluation metric.
For
Besides, more experimental details are revealed in Appendix We also discuss additional experi- ments in Appendix [ including solutions to regression problem and model stability given different spurious strengths at train time.
4.1 Simple Datasets with MLP
This section includes two simple datasets, Adult-Confounded and CMNIST, where spurious corre- lations are synthetically produced with the predefined strengths. For all competitors, MLP taken as the base model and full-batch training is implemented. Since KerHRM performs unstably over random seeds, we first average the results after 10 runs as its first score, and select top 5 among them as the second one, which will be marked with an asterisk (*) in each table.
4.1.1 Discussions on Adult-Confounded
We take UCI Adult 16 to predict binarized income levels (above or below $50,000 per year) Following [6], individuals are re-sampled according to sensitive features race and sex to simulate spurious correlations. Specifically, with binarized race Blacck/Non-Blacke and sex (Femalc/Male), four possible subgroups are constructed: Non-black Males (SGI), Non-black Females (SG2), Black Males (SG3), and Black Females (SG4). Keeping original train/test split and subgroup sizes from
UCI Adult dataset is widely used in algorithmic fairness papers. However, recent study g discusses some limitations of this dataset, such the choice of income threshold.

---

Table 2: P(Y 1|SG) for Adult- Confounded. IID shares spurious cor- relations with the train set. IND has no bias on race and sex. OOD defines the worst-case performance.
Table 3: Testing accuracy (%) on Adult Con-
founded. Three subsets are defined in Table 2 8 = 0 and 7 = 0 indicate the removal of Iu and LP when training ME-
IID
IND
OOD
Test (IID)
Test (IND)
Test
Train
ERM
92.4 0.1
66.8 0
40.7 0.5
(OOD)
BIIL
76.2 0.
73.5 0.5
70.2 #:
SGI
0.9
0.9
0.5
0.1
KerHRM
82.4 39
75.1 1:40
67.9 9.3
SG2
0.1
0.1
0.5
0.9
KerHRM
81.2
78.5 0:03
75.6
SG3
0.9
0.9
0.5
0.1
EDNIL
80.7 0.
79.1 0.4
77.5 =
SG4
0.1
0.1
0.5
0.9
EDNIL
91.8 =00
66.7 a
41.3 0.7
EDNIL
78.2
75.4 1.6
72.5 +
IRM
79.9 0.4
79.3 0.
78.8 0.4
UCI Adult, we sample data based on the given target distributions each sensitive subgroup as shown in Table [ In particular, OOD contributes the worst-case performance to validate if the predictions rely on group information. In this task, MLP with one hidden layer 96 neurons considered. For IRM, four environments comprise &arale. where the correlations between variant features (race, sex)
and target Y are distributed without overlapping. More details are provided Appendix B
With strong spurious correlations at train time, ERM obtains high accuracy as the correlations remain aligned; however, its generalization to other testing distributions is limited. Among all invariant learning methods without prior environment labels, EDNIL can perfectly identify variant features and infer ideally diversified environments. Therefore, it achieves the most invariant performance over different testing distributions. EIIL can improve consistency to some degree, but not as strong as EDNIL. A possible reason is that empirically trained reference model is not guaranteed to be purely variant [6) For KerHRM, it performs inconsistently across random seeds, which is reflected in the large standard deviation. In some cases, the performance hardly improves over iterations, observed by Liu al. 22)
Results
The results are shown in Table 3
Ablation Study for LEI We first claim the importance LLi, which constraints label dependency. by setting the coefficient 3 zero. As discussed in Section 3 the resulting environments are determined purely by target labels, and thus lead to inferior performance for invariant learning as shown in Table Next, we demonstrate the effectiveness of joint optimization in Figure 3a The regularization LIP promotes environment inference, so that the worst-case performance improves and remains stable over iterations. According Table Bl if the coefficient is turned off, feedback generated by Ml. will be ignored and the effect of invariant learning risks being undesirable.
4.1.2 Discussions on CMNIST
We report evaluation noisy digit recognition dataset, CMNIST. Following [1
we first assign Y= = to those whose digits are smaller than 5 and Y 0 to the others. Next, we apply label noise by randomly flipping Y with probability 0.2. Finally, the digits are colored with color label C, which is generated by randomly flipping Y with probability e. For training, two equal-sized environments with e 0.15 on average. For
with e = 0.1 and e 0.2 merged, which is equivalent testing, three situations are considered when e is 0.1, 0.5 or 0.9, respectively. Note that when e = 0.1, the spurious correlation is much aligned with the training set. On the other hand, e 0.9 defines the most challenging case since the spurious correlation shifts most dramatically from training.
For all competitors except KerHRM, we select MLP with two hidden layers of 390 neurons, and consider the whole dataset (50,000 samples) for training. For KerHRM who requires massive computing resources,
follow the settings recommended by [23] Specificallyy 5,000 samples and train MLP with one hidden layer of 1024 neurons. To construct ideally diversified Evracle for IRM, we pack all examples with C Y into environments, and C & Y into the other.
randomly select
Results The results are shown in Table First of all, not surprisingly ERM still adopts poorly to distributional shifts. Among all invariant learning methods without manual environment labels,
7

---

(b)
(c)
(d)
Figure 3: Analysis results. (a) Ablation study of joint optimization on Adult-Confounded, where y =
O means the removal of Ly. (b) CMNIST Performance of EDNIL with different configured numbers
of environments. (c) Comparison between Eeaam inferred by EDNIL and the oracle environments Eoracle CMNIST. (d) Testing EIlL''s and EDNIL's sensitivity to initialization Waterbirds.
EDNIL gets closest to IRM with Eoracles achieving consistent and robust performance in this dataset.
As shown in Figure 3c EDNIL provides almost ideally diversified Eleam for invariant learning.
Number of Environments As shown in Fig- 36 EDNIL is not sensitive to the predefined number of environments. Specifically, when the environment number is larger than the oracle (i.e. 2), some environment classifiers become re- dundant. Each of them provides moderate stant loss, taking up fixed and ignorable space in the softmax function. The visualization of Elear with 5 available environments is shown in Figure 3c] Additionally, training Ma in EDNIL with more environments is much more efficient than clustering based methods, such as the one proposed in KerHRM..
Table 4: Testing accuracy (%) of CMNIST, where color noise 0.9 makes the worst-case environment.
Color Noise
0.1
0.5
0.9
ERM
88.4 0.3
55.0 a.
21.7 # 0.8
EII.
79.6 0.3
71.7 07
63.1 = 0.5
KerHRM
74.3 0.7
66.2 + 1.7
58.0 + 11.5
KerHRM
71.3 0.7
68.5 4:05
66.1 # 0.7
EDNIL
77.7
76.8
75.2
IRM
77.8
76.8 0:04
75.2 = 0.3
4.2 Complex Datasets with Pre-trained Deep Learning Models
This section extends MLP to deep learning models with pre-trained weights for more complex data. With mini-batch fine-tuning, we consider all competitors but exclude KerHRM due to the efficiency issue. In Section 4.2.1 image dataset, Waterbirds 20 with controlled spurious correlations is selected for evaluating the generalization on more nigh-dimensional images. In Section 4.2.2 real-world NLP dataset, SNLI [3) considered. The biases in SNLI are naturally derived from the procedure of data collection, and we define biased subsets for evaluation following Dranker et al. [9]
4.2.1 Discussions Waterbirds
In Waterbirds 29 each bird photograph, from CUB dataset B1 is combined with one background image, from Places dataset 33 Both birds and backgrounds are either from land or water, and our target to predict the binarized species of birds. At train time, landbirds and waterbirds frequently present in land and water backgrounds respectively. Therefore, empirically trained models are prone
learn context features, and fail to generalize as background varies 6 IC 21 29)
To split a validation set whose overall distribution is i.i.d. to the training set, we merge original training and validation data Fl and split 10% for hyper-parameter tuning. For testing, observe all four combinations of birds and backgrounds in the original testing set. Among them, the minor subgroup (waterbirds on land) contributes the most challenging case. In this task, Resnet-34 [14] is chosen for mini batch fine-tuning. Given two binary labels (target, background), we distribute target background and target background into two different environments and apply balanced class weights for the oracle settings of IRM.
The results shown in Table As observed in 6 29 (i.e. waterbirds on land). EIIL also performs poorly in this case. With a more sophisticated learning
Results
ERM suffers in the hardest
'In the original training split, backgrounds are unequally distributed in each class, However, in the original validation split, they equally distributed, which is not i..dd to the training.

---

Table Testing accuracy (%) Waterbirds, where Y and BG means target and background respectively. The subgroup (Water, Land) contributes the worst-case performance.
(Y, BG)
(Land, Land)
(Water, Water)
(Land, Water)
(Water, Land)
ERM
99.4 0.0
91.4 0.2
90.9
72.8 1.0
EIIL
99.4 0.3
90.5 188
89.3 3.
68.6 5.54
EDNIL
98.5 0.
89.9 1.5
90.3
78.6 4:43
IRM 98.0 0.
90.6 1.1
89.5
83.2 2.2
framework, EDNIL narrows the gaps between subgroups and raises the worst-case performance. The results show that EDNIL is more resistant to distributional shifts.
Choice of Initialization Both EIll. and EDNIL take ERM as initialization. As mentioned in Section I. heavy dependency initialization is risky when testing distribution is unavailable. Therefore, we take ERM with different training steps for EIIL and EDNIL to verify the stability. The results are shown in Figure 3d As suggested by (6. EIIL works only with underfitted reference model in this case. If the reference model is more well-trained, the performance of EIIL will greatly decline since ERM might get distracted from variant features. One prone to be misled into an undesirable choice for EIIL when seeking syper-paramete instance, the validation score of EIIL with 500-step reference model (95.8%) ()) is higher than that with 100-step (94.7%), which is not consistent with their performances on testing. In comparison, EDNIL performs more consistent across different pre-training steps, which accentuates our strength of less sensitive initialization.
without prior knowledge of distributional shifts. For
4.2.2 Discussions on SNLI
The target of SNLI 3 is to predict the relation between two given sentences, premise and hypothesis. reveal hypothesis bias in SNLI, which is characterized by patterns in hypothesis sentences highly correlated with specific label. One can achieve low empirical risk without considering premises during prediction. However, as the bias no longer holds, the performance degradation occurs ||] 24
Recent studies 12
24 2/
We sample 100,000 examples and consider all classes, entailment, neutral and contrudiction, for experiment. Following ( we define three subsets, unbiased, bias aligned and bias misaligned, by training biased model with hypothesis as its only input. The specification of the subsets is as follows:
Unbiased: Examples whose predictions from the biased model are ambiguous
Aligned: Examples that the biased model can predict correctly with high confidence
Misaligned: Examples that the biased model can predict incorrectly with high confidence
The proportions of the three subsets are 17%, 67% and 16%, respectively. Due to the minority, the bias misaligned subset is more likely to be ignored and thus defines the worst-case performance.
For all methods, DistilBERT 50 taken as the pre-trained model for further mini-batch fine- tuning. For Eoractes we assign the bias aligned subset to the first environment, and the bias mis- aligned subset to the second. In order to make bias prevalence equal in the two environments. unbiased samples are scattered proportionally to the two environments.
Table 6: Testing accuracy (%) on SNLI, where the misaligned defines the worst-case performance.
Subset
Unbiased
Aligned
Misaligned
ERM
74.6 = n.
94.7 02
52.6 + 09
EIIL
74.2 :: 0.3
95.0 :: 0.1
51.7 :: L.3
EDNIL
74.3 0.
94.2 0.2
54.5 1.0
Results
The results are shown in Table 6 As
IRM 74.0 0.
92.3 0.5
56.9 + 1.1
reported by (9. ERM receives higher score on the bias aligned subset, but fails the bias misaligned case. Among all invariant learning methods without environment labels, only EDNIL improves on the bias misaligned subset. Namely, even though the definitions of biases are at high level, EDNII. is still capable of encoding and diversifying possible variant features.

---

Limitations
Our learning algorithm for environment inference is based the graphical model plotted in Figure As data are not necessarily generated by the presumed process, there can exist biases that cannot be captured by the proposed neural network. the paper, we provide empirical studies of effectiveness on diverse datasets, while we are still aware that a stronger guarantee of performance is required.
6 Conclusions and Societal Impacts
This work proposes EDNIL for training models invariant to distributional shifts. infer environments without supervision, propose multi-head neural network structure identify and diversify plausible environments. With joint optimization, the resulting invariant models are shown to more robust than existing solutions data having distinet characteristics and different strengths of biases. We attribute the effectiveness to the underlying learning objectives, which are consistent with recent studies of ideal environments. Additionally, note that the classifier-like structure of environment inference model makes EDNIL easy to combine with off-the-shelf pre-trained models and trained more efficiently.
Our contributions to invariant learning have broader societal impacts on numerous domains. For instance, it can encourage further research and real-world applications debiasing machine learning systems. To be more specific, the identification and elimination of potential biases can facilitate more robust model training. It can be beneficial to many real applications where distributional shifts commonly occur, such as autonomous driving, social media and healthcare.
Furthermore, discussed by (6), invariant learning can promote algorithmic fairness in some ways. In particular, our empirical achievements on Adult-Confounded can prevent discrimination against sensitive demographic subgroups in decision-making process. It shows that EDNIL has potential to
learn a fair predictor without prior knowledge of sensitive attributes, which related to [3 19 We expect that one can extend our work to more fairness benchmarks and criteria in the future.
Last but least, it is worth mentioning some cautions, however. Since the invariant learning algorithm claims to find invariant relationships, one might cast more attention on feature importance of the invariant model and even incorporate the results into further research or applications. Nevertheless. the results are reliable only when the model is trained appropriately. Insufficient data collection or careless training process, for example, can certainly affect the identification of invariant features, and thus mislead experimental findings. As result, we believe that adequate and careful preparations and analyses are essential before drawing conclusions from the inferred invariant relationships.
Acknowledgments and Disclosure of Funding
We would like to thank the anonymous reviewers for their helpful suggestions. This material is based upon work supported by National Science and Technology Council, ROC under grant number 111-2221-BAA 2--2022------ -146 -MY3 and 1 110-6644-F-0022
10

---

References
[1] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk mini- mization, 2019. URL http:attt::/aaii..orr///a/ http//.119700093333
[2] Sara Beery, Grant Van Hom, and Pietro Perona. Recognition in terra incognita. In Proceedings of the European Conference on Computer Vision (ECCV), September 2018.
[3] Samuel R. Bowman, Gabor Angeli, Christopher Potts, and Christopher D. Manning, A large annotated corpus for learning natural language inference. In Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing, pages 632-642, Lisbon, Portugal, sep 2015. Association for Computational Linguistics. doi: 10. 10.18653/1/DDI5-107. URL
httpe://acclamthologgy.. 101.1//100..0/11--175555 1075
[4] Shiyu Chang. Yang Zhang, Mo Yu, and Tommi Jaakkola. Invariant rationalization. In Proceed- ings of the 37th International Conference on Machine Learning, pages 1448-1458, 2020. URL
https::/proceeddnngg mlr ppess/vv19) chang20c. hhmml
[5] Yo Joong Choe, Jiyeon Ham, and Kyubyong Park. An empirical study of invariant risk minimization. In ICML 2020 Workshop on Uncertainty and Robustness in Deep Learning. 2020.
[6] Elliot Creager, Joem-Heerik Jacobsen, and Richard Zemel. Environment inference for invariant learning. In Proceedings of the 38th International Conference on Machine Learning, pages 2189-2200,, 2021. URL https: http::://prooceeiine
mlr. wwwpssreesse creager httl
[7] Jacob Devlin, Ming-Wos Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Confer- ence of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pages 4171 Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10. 10.1863//1/N19-1444 14233.
URL httppsss//aaaaaaoollggoooooo/eee orggoogg.oor//// N11//11991111 1423
[8] Frances Ding, Moritz Hardt, John Miller, and Ludwig Schmidt. Retiring adult: New datasets for fair machine learning. In Advances in Neural Information Processing Systems, pages 6478-6490, 2021. URL https:: httttps:/ppooeeedinnngg
32e5444
pdf
[9] Yana Dranker, He He, and Yonatan Belinkov. IRM case of natural language inference. In Advances in Neural Information Processing Systems,
when it works and when it doesn't: A test
2021. URL http:ttts///operrrriieeee wttp:/oor////ddttttttttob..
[10] Robert Geirhos, Jörn-Henrik Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Bren- del, Matthias Bethge, and Felix A. Wichmann. Shortcut learning in deep neural networks. Nature Machine Intelligence, 2(11):665 673, 2020. URL hctps://ddi.oor/1000 1038%
[11] Max Glockner, Vered Shwartz, and Yoav Goldberg. Breaking NLI systems with sentences that require simple lexical inferences. In Proceedings of the 56th Annual Meeting of the Association for Computationnal Linguistics ( (Volume 2: Short Pupers), pages 650-655, Melbourne, Australia, July 2018. Association for Computational Linguistics. doi: 10. 10.18853/ URL https://aclanthology wtt////-//110
[12] Suchin Gururangan, Swabha Swayamdipta, Omer Levy, Roy Schwartz, Samuel Bowman, and Noah A. Smith. Annotation artifacts in natural language inference data. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 2 (Short Papers), pages 107-112, June 2018. doi: 10.18653/v 1LL16653/1/N18-201111 URL https:://aclanttoooogoooooo// org/ 1118-2117
[13] Tatsunori Hashimoto, Megha Srivastava, Hongseok Namkoong, and Percy Liang. Fairness without demographics in repeated loss minimization. In Proceedings of the 35th International Conference on Machine Learning, pages 1929-1938, 2018. URL htRls:ttps:/// htattp:://proceeding.
mlr.

---

[14] K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 770-778, Los Alamitos, CA, USA, jun 2016. IEEE Computer Society. doi: 10.1 101109/CVPR.0016 URL https://doi.. www.comttterssscstttt..... wrg.://iiiiii..o.o.ooooo 11t9//..1111/P...1111 2016.90
[15] Arthur Jacot, Clément Hongler, and Franck Gabriel. Neural tangent kemnel: Convergence and generalization in neural networks. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Resa-Biaanci, and Roman Garnett, editors, Advances in Neural Information Processing Systems 31: Anmual Conference Neural Information Processing Systems 2018, NeurlPS 2018, 3-8 December 2018, Montréal, Canada, pages 8580-8589, 2018.
[16] Ronny Kohavi and Barry Becker. Adult. UCI Machine Learning Repository, 1996. Accessed: 2021-09-2 109000
[17] David Krueger, Ethan Caballero, Zhang. Remi Le Priol, and Aaron Courville. Out-of-distributie generalization via risk extrapo- lation (rex). In Proceedings of the 38th International Conference on Machine Learning, pages 5815-5826, 2021. URL httpp:://pooceeding. mlr wwwpppresssvveeee wv19s//vvvv//////
Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai
(//////zz 1a111399//ruee222aa htm
[18] Kun Kuang, Ruoxuan Xiong. Peng Cui, Susan Athey, and Bo Li. Stable prediction with model misspecification and agnostic distribution shift. In Proceedings of the AAAI Conference on Artificial Intelligence, pages 1485-4492) 04 2020. doi: 10. 16099/aai.v34404.5555..
[19] Preethi Lahoti, Alex Beutel, Jilin Chen, Kang Lee, Flavien Prost, Nithum Thain, Xuezhi Wang, and Ed Chi. Fairness without demographics through adversari- ally reweighted learning. In Advances in Neural Information Processing Systems, pages 728-740, 2020. URL htttps://proceeddnggg www.snnuuriisiiceeoe potrrpppps..c/ppppr////////
pdf
[20] Yong Lin, Shengyu Zhu, and Peng Cui. Zin: When and how to learn invariance by environment inference?, 2022. URL URL httttt////axii wwwpxxii..rrvaaaee org/ abs/ abs//b03.... 05818.
[21] Evan Z Liu, Behzad Haghgoo, Annie S Chen, Aditi Raghunathan, Pang Wei Koh, Shiori Sagawa, Percy Liang, and Chelsea Finn. Just train twice: Improving group robustness without training group information. In Proceedings of the 38th International Conference on Machine Learning, .ttll
pages 6781- 2021. URL https: http:///pocceeeeeeeeeeoee mlr press/vv39/Iise
[22] Jiashuo Liu, Zheyuan Hu, Peng Cui, Bo Li, and Zheyan Shen. Heterogeneous risk minimization. In Proceedings of the 38th International Conference Machine Learning, pages 1804-6814 2021. URL httpshttps:/// http:::///coccepepeeeeeeee mlr wwwpseeeiiiiee hwl////119//111 html
[23] Jiashuo Liu, Zheyuan Hu, Peng Cui, Bo Li, and Zheyan Shen. Kernelized hetero- geneous risk minimization. In Advances in Neural Information Processing Systems. pages 21720-21731 2021. URLLhttps://proceedii www.ddiinnnnnnnnnnioeeeeee 155a5118000ff6622888888144888881
Paper. pdf
[24] Tom McCoy, Ellie Pavlick, and Tal Linzen. Right for the wrong reasons: Diagnosing syn- tactic heuristics in natural language inference. In Proceedings of the 57th Annual Meet- ing of the Association for Computational Linguistics, pages 2428-4448, Florence, Italy, July 2019, Association for Computational Linguistics. doi: 10. 10.188 URL https://allatthologg.ooo
1334
[25] Judea Pearl. Probabilistic Reasoning in Intelligent Systems. Morgan Kaufmann, 1988.
[26] J. Peters, Peter Buhlmann, and Nicolai Meinshausen. Causal inference by using invariant prediction:
identification and confidence intervals. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 78, 2015.
[27] Adam Poliak, Jason Naradowsky, Aparajita Haldar, Rachel Rudinger, and Benjamin Van Durme. Hypothesis only baselines in natural language inference. Proceedings of the Seventh Joint Conference on Lexical and Computational Semantics, pages 180-191, New Orleans, Louisiana, jun 2018 Association for Computational Linguistics. doi: 10. 1/5518/611111222000 URL.
org/ SS82
12

---

[28] Mateo Carulla, Bernhard Schölkopf, Richard E. Turner, and J. Peters. Invariant models for causal transfer learning. J. Mach. Learn. Res., 19:36:1-36:34, 2018.
[29] Shiori Sagawa*, Pang Wei Koh*, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust neural networks. In International Conference Learning Representations, 2020. URL
https://openreries:
[30] Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. In NeurlPS 2019 Workshop Energy Efficient Machine Learning and Cognitive Compuling, 2019.
[31] P Welinder, S Branson, T Mita, C Wah, Schroff, S Belongie, and P Perona. Caltech-UCSD Birds 200. Technical Report CNS-TR-2010-001, Caltech, 2010.
[32] Shen Yan, Hsien-te Kao, and Emilio Ferrara. Fair class balancing: Enhancing model fairness without observing sensitive attributes. In Proceedings of the 29th ACM International Conference on Information Knowledge Management, CIKM '20, page 1715-1724, New York, NY, USA,
2020. Association for Computing Machinery. ISBN 978 978145036,,9, doi: 10.1145/3340511. 3411980. URL https:://oi..... org/ 10. 1145/3340511.... 3411980
[33] Bolei Zhou, Agata Lapedriza, Aditya Khosla, Aude Oliva, and Antonio Torralba. Places: A
10 million image database for recognition. IEEE Transactions Pattern Analysis and Machine Intelligence, 40(6):1452-146; 2018. doi: 10.1
13


Reviews:
[{"review_id": "f4iHIsm-mGk", "reviewer": "Reviewer_T3xK", "summary": "The authors propose a novel environment-free invariant learning method that uses an auxiliary network to learn environment-specific features, from which environment inferences can be derived. These inferred environments labels are then used to train and invariant model.\n\nUPDATE 2022-aug-05: The paper's clarity has been improved so I am increasing my score by one.", "questions": "* I was a little unclear on the motivation for using the environment (in)dependency objective - I can see its utility in the ablation studies, but what was the reason, conceptually speaking, to introduce this in the first place?\n* Beyond looking at the value of L_ED during training, can an experiment be formulated to test whether the inferred environments (samples from P(e|X,Y)) are indeed “diverse” in the way the authors are hoping for?\n* Section 3 is a bit hard to parse given that the three objectives are introduced in several places (lines 129–133, Fig 1, Eqn 6). I would suggest that the authors take care to discuss these in the same order each time and include the symbols alongside the objective names (e.g. “Label dependency objective” => “Label dependency objective ($L_{LD}$” in line 129). I also found it a bit confusing that some objectives are for the inference net Phi while others are for the invariant learner. While I can see that Fig 1 attempts to make this distinction by placing the loss terms in different parts of the figure, it would be even clearer if the networks (or their parameters) were used as arguments to the losses, e.g. L_ED(Phi, Psi) = L_ED(Psi) + beta L_LD(Psi) + gamma L_IP(Phi) in Eqn 6). I would also order the terms in Eqn 6 in the same order as they appear in the text.\n* Since the variance of per-env losses is used as L_IP, it is worth mentioning REx (https://arxiv.org/abs/2003.00688) as a related work, since it uses a similar objective (with known environments).\n* Regarding the tabular experiments, there are some limitations of the UCI-Adult dataset (https://openreview.net/forum?id=bYi_2708mKK) that would be relevant/interesting to NeurIPS readers/attendees (perhaps in a footnote).\n* Environment-free invariant learning is a very ambitious problem, which could be difficult to solve in some settings. I would be curious to hear about experimental settings where EDNIL fails; I think these would interest the reader and also make the limitations of this method a bit more clear.\n* When mentioning how the training/validation splits for Waterbirds are merged [lines 290–291], it would be good to mention that the original validation split contains all four groups in equal proportions, whereas original training split has unequal proportions.\n* I find the double negatives needed to parse Table 1 to be slightly confusing.\n", "limitations": "No [line 430]. I suggest that the authors add a discussion of possible societal impacts to a future revision. Given the connections of invariant learning to fairness (https://arxiv.org/abs/2010.07249) and the use of a fairness benchmark in the experiments, there should be plenty to discuss.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "I like the idea of the paper, and was impressed by the experimental results showing that the proposed method EDNIL realizes good worst-group performance in a variety of settings, while addressing the sensitivity to initialization that we see in EIIL. In my opinion, the experimental results are sufficient but the presentation should be improved (see below for specific suggestions).\n\nStrengths\n* I like the experimental results. The “effortless initialization” result (Fig 7), showing that the proposed method is not sensitive to initial ERM features, which is a known issue with EIIL, is especially impressive.\n* The paper does a good job of describing related approaches to environment-free invariant learning, and comparing against these methods empirically.\nWeaknesses\n* The paper can be hard to read at times. Some parts would benefit from additional exposition. For example the idea of “diversity” of inferred environments is mentioned in passing with a citation, but it is not fully explained what type of diversity is needed, and how it can be measured formally? The cited paper (https://arxiv.org/abs/2004.05007) mentions that the correlation pattern between Y and spurious features should vary across environments…is the outcome we hope for when introducing L_ED?\n", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "YsbZdE5IgTi", "reviewer": "Reviewer_8Cwm", "summary": "\nThe paper discusses the problem of shift between train and test data distributions. The authors assume that training data originate from multiple 'environments', each with particular distribution. They propose a new framework (EDNIL) through which they can infer the environments from data (without human labeling) and then train an environment invariant model. The principal assumption is that of the existence of `invariant features` - the distribution of the outputs conditioned on these is the same across all the environments. The loss is a combination of multiple terms motivated through information theoric arguments. The algorithm consists of alternating steps to infer the environments and train the invariant model. Multiple experiments are provided to corroborate the effectiveness of the method.", "questions": "1) While the intro speaks about a shift between train and test distributions, in the end it seems you address the problem of a mix of multiple distributions in both train and test. That is the train and test are not fundamentally different, they are just both formed through a combination of multiple distributions, is this right?\n2) If so, what is the meaning of $\\mathcal{E}_{all}$ in your setup?\n3) Do I assume correctly that the output and input spaces are aligned across the environments (including the supports of the distributions)?\n4) What are $\\Phi$ and $\\Psi$? Are these some general functions such as small MLPs?\n5) In equation (5) softmax you sum across the E_learn invoronments, Does it mean that the number of environments needs to be fixed beforehand?\n6) How is the environment decided at test time, when no labels are available? Which head shall be used? ", "limitations": "Societal impact - not addressed, not directly relevant.\nLimitations - invalidity of graphical model assumption is mentioned. Yes, is reasonable. ", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Train/test distribution shift is an important and fully answered problem, particularly relevant for all practical applications where the standard iid assumptions typically do not hold. As such the paper addresses a significant question for the community.\nI found the paper somewhat difficult to parse and follow. Most importantly, I found the description of the whole setup (shift from train to test distributions or a mix of multiple environments in train and test) rather confusing. But I have more question marks (see Questions) indicating a fairly low clarity of the paper. \nIn the end the solution is an alternating optimization in which through the first step the environments are inferred though a form of supervised clustering (similarity metric is the model accuracy) and then this info is used to train a prediction model. If my interpretation is correct, this does not seem very novel.", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "FnWpxXNQQG6", "reviewer": "Reviewer_SNzy", "summary": "Paper proposes a new method to infer environment labels that can then be used for invariant learning (such as IRM methods). Essentially, it uses a neural network for clustering data into environments with appropriate loss functions. The method is more scalable and less sensitive to initializations than the prior methods. Experiments show significant improvements for distribution-shifted tasks without environment labels.", "questions": "1. Problem 2.1 needs to be explained more, discussing the use of Shannon entropy. Also clarify the novelty; I believe Equation (3) was proposed in [1] but Equation (4) is a contribution?\n2. In Equation (4), is it enough that LHS is just > 0 while being arbitrarily small? Should the LHS be maximized instead?\n3. The proposed method in Equation (5) needs more explanation on why it was chosen to be this way. I believe it is done to ensure that data is divided into environments that maximize label prediction in the individual environments. This is interesting but requires more justification in the text. Related to my previous comment, $H(Y | X_v) - H(Y | X_v, \\mathcal{E}_\\text{learn})$ might be actually maximized with this objective, and not just >0. \n\n4. Discussing $L_{IP}$ before discussing the $M_{IL}$ model hampers clarity. It should also be clarified/reiterated that $L_{IP}$ is only used to update the $M_{EI}$ network and not $\\Phi$, and that there are two separate invariance preserving objectives (one for $M_{IL}$ and one for $M_{EI}$). \n5. It is very interesting that the number of environments (clusters) need not be specified accurately for EDNIL to work (as seen in Figure 6). More discussion regarding why this is so will be helpful, i.e., which part of the loss ensures this.\n\n\nMinor:\n\n1. Please use \\text{} in math equations for non-variables, such as \\text{learn}, \\text{LD}, etc. Also, use \\exp in Equation (5).\n2. Line 19: “… aims at learning causality expected to be robust …” should be “..aims at learning causal features expected to be robust..”.\n3. Line 106: should be “HRM equips a joint learning framework”. \n\n[1] Jiashuo Liu, Zheyuan Hu, Peng Cui, Bo Li, and Zheyan Shen. Heterogeneous risk minimization. In International Conference on Machine Learning (ICML), 2021.", "limitations": "Yes", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. Paper tackles an important challenge of robust learning with no environment labels and solves the problems with existing works.\n\n2. Experiments in varied tasks, including images and text, show significant performance improvements. Experiments also help answer questions about the effect of initializations, effect of number of environments.\n\n\nWeaknesses:\n\n1. Clarity can be improved a lot; especially Section 2.2 and Section 3 (Methodology). As a result, the paper’s contributions are muddled with existing work (for example, which of the learning objectives and loss terms are novel?). \n2. Paper can be improved a lot with better justifications for the choices made in the paper. Currently the loss functions seem a little arbitrary; for example, why are Equations (5) and (9) the best ways to achieve their respective goals. \n3. There are no theoretical guarantees provided by the prior works, but this can be future work. \n\n", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "0GBvTERfUgQ", "reviewer": "Reviewer_QRiV", "summary": "In the present paper, the authors tackle the challenging problem of learning a supervised ML model capable of identifying and adapting to different environments hidden in the training data. The method, EDNIL, is composed of two jointly learned models, that take care of the environment identification, the learning of the invariant representations and the label predictions, produced by a multi-headed neural network. The proposed model is compared to different alternative models from the literature of the field, in different challenging benchmarks, and the results show that it closely achieves the best possible invariant learning performance. \n\nAFTER REBUTTAL\nThe authors have addressed the main concerns raised during the review. I confirm my assessment of this paper as solid and deserving of acceptance in the conference.", "questions": "- In general, the presence of way too many acronyms can confuse the reader. I do understand that presenting all this material in 9 pages is quite a challenge but the density is a bit excessive in some passages of the manuscript.\n- lines 74-93 seem a bit hard to follow, especially about the role of Phi and Xv.\n- the role of beta and gamma seems important, but the authors only give space in the main text to the discussion of the extreme cases beta=0 and gamma=0. Can a few comments be added about the range of optimal values for these parameters? How hard is it to fine-tune them and how variable are they across different learning problems?\n- lines 166-169 (and later figure 7 and its description): what is the intuitive explanation of the reduced dependence on the training stage for the ERM initialization. Is it about the role of the L_IP term, that avoids trivial environment split scenarios?\n- line 195 is not clear, probably just needs rewriting.\n- lines 227-230 the sentence seems to be missing a verb and becomes a bit obscure.\n- line 236: removing the coefficient beta -> in what sense? setting it to 0, or to 1? Probably rephrasing could help.\n- In general, it seems to me that this work could be very impactful also for the field of ML fairness, where the different environments could represent protected categories that are misrepresented in the available data and the learning model risks to become biased. Did the authors consider this connection?\n", "limitations": "I think the main limitations of the paper are well acknowledged by the authors.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The authors introduce a very flexible and effective learning framework, incorporating the main strengths of the previously introduced models and surpassing many of the limitations thereof. The novel method is shown to perform extremely well when compared to its competitors, achieving close to the performance of IRM with access to the oracle environments. \n\nI think the main weakness of this paper is in the density of the presentation and in the abundance of acronyms that make the read quite challenging, especially for a non-expert.", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept