Environment Diversification with Multi-head Neural Network for Invariant
  Learning
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 invariant 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.
Bo-Wei Huang, Keng-Te Liao, Chang-Sheng Kao, Shou-De Lin
# Environment Diversification with Multi-head Neural Network for Invariant 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 invariant 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.

## 1 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 (i.i.d.). However, data in real-world applications can come with undesired biases, causing a 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 [10]. In this work, we focus on _Invariant Learning_, which aims at learning invariant features expected to be robust against distributional changes. Invariant Risk Minimization (IRM) [1] has been proposed as a popular solution for invariant learning. 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 as a standard ERM loss function with a penalty term constraining the trained model (e.g. classifier) to be optimal in all the environments.

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 are summarized as follows.

Prior Knowledge of EnvironmentsIRM assumes training data are collected from environments, and the environment labels (i.e. which data instance belongs to which environment) are given. However, the environment labels are often unavailable. Moreover, the definition of environments can be implicit, making human labeling more difficult and expensive. To find environments without supervision, Creager et al. [6] propose Environment Inference for Invariant Learning (EIIL), a min-max optimization framework training the model of interest and inferring the environment labels. Another work, Heterogeneous Risk Minimization (HRM) [22], or its extension called Kernelized Heterogeneous Risk Minimization (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 on carefully chosen features satisfying a series of theoretical constraints, and thus human efforts are still required.

**Delicate Initialization**EIIL is able to infer environments but requires an ERM model for initialization. Crucially, the ERM model should heavily depend on spurious correlations. Creager et al. [6] 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 on raw features. As for KerHRM, although it extends HRM to avoid the issue of representation learning 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 a novel framework, **E**nvironment **D**iversification with multi-head neural **N**etwork for **I**nvariant **L**earning (EDNIL). EDNIL is 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 recent studies [5; 22; 23] 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 are summarized as:

* We implement this framework using various pre-trained models such as Resnet [14] and DistilBert [30], and evaluate it with diverse data types and varied biases. The results show that EDNIL can constantly outperform the existing state-of-the-art methods.
* The learning algorithm of EDNIL has theoretical connections to recent studies [5; 20; 22; 23] discussing conditions of ideal environments.
* EDNIL mitigates the three limitations discussed above. The comparisons between EDNIL and other methods are shown in Table 1.

## 2 Preliminaries and Related Work

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, we discuss recent studies investigating ideal environments. In Section 2.3, we introduce the existing unsupervised methods inferring environments.

### Out-of-Distribution Generalization and Invariant Learning

Following [1; 22], we consider a dataset \(D=\{D^{e}\}_{e\in\mathrm{supp}}(\mathcal{E}_{\texttt{v}})\) with different sources \(D^{e}=\{(x_{\texttt{v}}^{e},y_{\texttt{v}}^{e})\}_{i=1}^{n_{e}}\) collected under multiple training environments \(e\in\mathrm{supp}(\mathcal{E}_{\texttt{v}})\). Random variable \(\mathcal{E}_{\texttt{tr}}\) indicates training environments. For simplicity, \(X^{e}\), \(Y^{e}\) and \(P^{e}\) denote features, target and distribution in environment \(e\), respectively.

With \(\mathcal{E}_{\texttt{all}}\) containing all possible environments such that \(\mathrm{supp}(\mathcal{E}_{\texttt{all}})\supset\mathrm{supp}(\mathcal{E}_{ \texttt{tr}})\), the goal of out-of-distribution generalization is to learn a predictor \(f(\cdot):\mathcal{X}\rightarrow\mathcal{Y}\) as Equation 1, where

\begin{table}
\begin{tabular}{c c c c}  & Unsupervised1  & Insensitive Initialization & Efficiency \\ \hline IRM [1] & ✗ & ✓ & ✓ \\ ZIN [20] & ✗ & ✓ & ✓ \\ EIIL [6] & ✓ & ✗ & ✓ \\ HRM [22] & ✓ & ✓ & ✗ \\ KerHRM [23] & ✓ & ✓ & ✗ \\
**EDNIL (Ours)** & ✓ & ✓ & ✓ \\ \end{tabular}
\end{table}
Table 1: A summary of the advantages of invariant learning methods.
\(\mathbb{E}_{X^{\text{c}},Y^{e}}[l(f(X^{\text{c}}),Y^{e})]=\mathbb{E}^{e}[l(f(X^{ \text{c}}),Y^{e})]\) measures the risk under environment \(e\) with any loss function \(l(\cdot,\cdot)\). In general, for \(e\in\operatorname{supp}(\mathcal{E}_{\text{tr}})\) and \(e^{\prime}\in\operatorname{supp}(\mathcal{E}_{\text{all}})\setminus \operatorname{supp}(\mathcal{E}_{\text{tr}})\), \(P^{e^{\prime}}(X,Y)\) is rather different from \(P^{e}(X,Y)\).

\[f=\arg\min_{f}\max_{e\in\operatorname{supp}(\mathcal{E}_{\text{all}})}R^{e}(f) \tag{1}\]

Recently, several studies [1; 4; 18; 26; 28] have attempted to tackle the generalization problems by discovering invariant relationships across all environments. A commonly proposed assumption is the existence of _invariant features_\(X_{\text{c}}\) and _variant features_\(X_{\text{v}}\). Specifically, raw features \(X\) are assumed to be composed of \(X_{\text{c}}\) and \(X_{\text{v}}\), or \(X=h(X_{\text{c}},X_{\text{v}})\) where \(h(\cdot)\) is a transformation function. Invariant features \(X_{\text{c}}\) are assumed to be equally informative for predicting targets \(Y\) across environments \(e\in\operatorname{supp}(\mathcal{E}_{\text{all}})\). On the contrary, the distribution \(P^{e}(Y|X_{\text{v}})\) can arbitrarily vary across \(e\). As a result, predictors depending on \(X_{\text{v}}\) can have unpredictable performance in unseen environments. In particular, the correlations between \(X_{\text{v}}\) and \(Y\) are known as spurious and unreliable.

To extract \(X_{\text{c}}\), IRM [1] assumes there is an encoder \(\Phi\) for obtaining representations \(\Phi(X)\approx X_{\text{c}}\). The encoder is trained with a regularization term enforcing simultaneous optimality of the predictor \(w\circ\Phi\) in training environments, where dummy classifier \(w=1.0\) is a fixed multiplier for the encoder outputs:

\[\sum_{e\in\operatorname{supp}(\mathcal{E}_{\text{tr}})}R^{e}(\Phi)+\lambda|| \nabla_{w|w=1.0}R^{e}(w\circ\Phi)||^{2} \tag{2}\]

As \(w\) is a dummy layer, the encoder \(\Phi\) is also regarded as a predictor.

### Ideal Environments

As \(\mathcal{E}_{\text{tr}}\) is unavailable or sub-optimal in most applications, learning to find appropriate environments (denoted by \(\mathcal{E}_{\text{learn}}\)) is attractive. However, the challenge is the lack of knowledge of valid environments. Recently, Lin et al. [20] have proposed Equation 3 and 4 as the conditions of ideal environments, where \(H\) 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 to collect and verify the additional information.

\[H(Y|X_{\text{c}})=H(Y|X_{\text{c}},\mathcal{E}_{\text{learn}}) \tag{3}\]

\[H(Y|X_{\text{v}})-H(Y|X_{\text{v}},\mathcal{E}_{\text{learn}})>0 \tag{4}\]

Particularly, Equation 4 can be implied by empirical studies [5] where diversity of environments is recognized as the key to obtaining effective IRM models. To be more precise, large discrepancy of spurious correlations, or \(P^{e}(Y|X_{\text{v}})\), between environments is favored. As the environments give a clear indication of distributional shifts, IRM can easily identify and eliminate variant features. Beyond IRM, HRM [22] and KerHRM [23] can also be viewed as optimizing diversity via clustering-based methods specifically.

### Unsupervised Environment Inference

Here we provide a 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 [6] proposes formulating invariant learning as a min-max optimization problem. Specifically, EIIL is composed of two objectives, _Environment Inference_ (\(EI\)) and _Invariant Learning_ (\(IL\)), where \(EI\) is optimized by maximizing the training penalty via labeling the data, and \(IL\) is optimized by minimizing the training loss given the data labeled by \(EI\). The two-stage framework bypasses the difficulty of defining environments; however, the training result heavily relies on the initialization of the \(EI\) optimization. Specifically, the initialization demands a strongly biased ERM reference model; otherwise, EIIL can have a significantly weaker performance.

Another method, HRM [22], proposes a clustering-based 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, a known issue of HRM is an 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. [23] propose KerHRM. The main idea is to adopt the Neural Tangent Kernel [15] method which transforms non-linear neural network training into a linear regression problem on the proposed Neural Tangent Features space. As a 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 an affordable option at the present stage.

## 3 Methodology

In this section, we propose a general framework to learn invariant model without manual environment labels. As shown in Figure 0(a), our proposed method consists of two models, \(M_{\text{EI}}\) and \(M_{\text{IL}}\). Given the pooled data \((X,Y)\), \(M_{\text{EI}}\) infers environments \(\mathcal{E}_{\text{learn}}\) satisfying Condition 3 and 4, and \(M_{\text{IL}}\) is an invariant model trained with the inferred environments. Our framework is jointly optimized with alternating updates. The learned \(M_{\text{IL}}\) can provide information of invariant features to \(M_{\text{EI}}\), so that Condition 3 and 4 can be fulfilled simultaneously. Note that \(M_{\text{EI}}\) only serves at train time to provide environments for invariant learning. At test time, only \(M_{\text{IL}}\) is needed to perform invariant predictions.

### The Environment Inference Model

The target of environment inference is to partition data into environments \(\mathcal{E}_{\text{learn}}\) satisfying Condition 3 and 4. In this regard, we propose a graphical model, which is a sufficient condition for Condition 3 and 4 (the proof is in Appendix A), 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 [1].

The inference model, \(M_{\text{EI}}\), is a neural network learning to realize the underlying graphical model. Following the idea of parameterizing environments from HRM [22] and KerHRM [23], we assume the distinct mapping relation between \(X\) and \(Y\) in environment \(e\) can be modeled by a function \(f^{e}(\Psi(X))\), where \(\Psi(X)\) is learned representations expected to encode variant features \(X_{\text{v}}\) and \(f^{e}\) is an environmental function responsible for predicting \(Y\). Instead of employing clusters, we propose building a multi-head neural network as shown in Figure 0(b); particularly, a single-head network

Figure 1: Concepts of EDNIL. (a) The joint learning framework of EDNIL. The environment inference model \(M_{\text{EI}}\), containing a variant encoder \(\Psi\) and environmental functions \(f^{e}\), is trained with \(L_{\text{EI}}=L_{\text{ED}}+\beta L_{\text{LI}}+\gamma L_{\text{IP}}\). The invariant learning model \(M_{\text{IL}}\), containing an invariant predictor \(\Phi\), is trained with \(L_{\text{IL}}\). (b) The multi-head network structure of the environment inference model \(M_{\text{EI}}\).

Figure 2: The graphical model.
in \(M_{\text{EI}}\) with shared parameters corresponds to a cluster center in HRM or KerHRM. The training procedure of \(M_{\text{EI}}\) can be divided into _inference stage_ and _learning stage_.

#### 3.1.1 Inference Stage of \(\boldsymbol{M_{\text{EI}}}\)

The goal is to infer an environment label for each training data. As in the graphical model, \(\mathcal{E}_{\text{learn}}\) is associated with variant relationships. Inspired by multi-class classification problem, we propose Equation 5, where the probability \(P(e\mid X,Y)\) is estimated via a softmax of negative \(l(f^{e}(\Psi(X)),Y)\) divided by a constant temperature \(\tau\). The function \(l\) is expected to be the commonly used cross entropy or mean squared error that measures the discrepancy between \(Y\) and \(f^{e}(\Psi(X))\) for each environment \(e\). Intuitively, each data prefers the environment whose model has better prediction.

\[P(\mathcal{E}_{\text{learn}}=e\mid X,Y)=\frac{\exp{(-l(f^{e}(\Psi(X)),Y)/\tau)} }{\sum\limits_{e^{\prime}\in\operatorname{supp}(\mathcal{E}_{\text{learn}})} \exp{(-l(f^{e^{\prime}}(\Psi(X)),Y)/\tau)}} \tag{5}\]

#### 3.1.2 Learning Stage of \(\boldsymbol{M_{\text{EI}}}\)

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.e. _Environment Diversification Loss_ (\(L_{\text{ED}}\)), _Label Independence Loss_ (\(L_{\text{LI}}\)) and _Invariance Preserving Loss_ (\(L_{\text{IP}}\)). In particular, \(L_{\text{ED}}\) and \(L_{\text{IP}}\) correspond to the concepts of Condition 4 and 3, respectively.

Environment Diversification Loss (\(L_{\text{ED}}\))We consider maximizing \(H(Y|X_{\text{v}})-H(Y|X_{\text{v}},\mathcal{E}_{\text{learn}})\) to satisfy Condition 4 and capture more diverse variant relationships. Given the estimated \(P(\mathcal{E}_{\text{learn}}|X,Y)\), \(L_{\text{ED}}\) selects the most probable environment and its corresponding network for optimization:

\[L_{\text{ED}}=-\sum_{i}w_{i}\max_{e}[\log P(e\mid x_{i},y_{i})] \tag{6}\]

For each data \((x_{i},y_{i})\), although only one environment \(e_{i}=\operatorname*{arg\,max}_{e}P(e\mid x_{i},y_{i})\) is selected for the minimization, the softmax simultaneously propagates gradient to maximize \(l(f^{e^{\prime}}(\Psi(x_{i})),y_{i})\) for \(e^{\prime}\neq e_{i}\). The network learns to maximize the dependency between \(\mathcal{E}_{\text{learn}}\) 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 \(w_{i}\) inversely proportional to the size of \(e_{i}\). The importance of smaller environments will be thus enhanced within the summation.

Label Independence Loss (\(L_{\text{LI}}\))With d-separation [25], \(\mathcal{E}_{\text{learn}}\) is independent of \(Y\) in the graphical model. Hence, \(L_{\text{LI}}\) constraints their dependency measured by mutual information \(I(Y;\mathcal{E}_{\text{learn}})\). It can be verified that minimizing \(I(Y;\mathcal{E}_{\text{learn}})\) is equivalent to minimizing Equation 7 given \(P(Y)\) is known. Empirically, \(L_{\text{LI}}\) prevents a trivial solution that environments are determined purely by target labels regardless of input features, which is undesirable for invariant learning.

\[L_{\text{LI}}=\mathbb{E}_{e\sim P(\mathcal{E}_{\text{learn}})}[\sum_{y}P(y|e) \log P(y|e)] \tag{7}\]

Invariance Preserving Loss (\(L_{\text{IP}}\))For Condition 3, as \(M_{\text{IL}}\) learns some invariant relationships after several training steps (Section 3.2), \(L_{\text{IP}}\) can be considered to exclude invariant features from the diversification. Specifically, designed as the contrary of \(L_{\text{ED}}\), \(L_{\text{IP}}\) limits the variance of expected loss from invariant predictor \(\Phi\) (in \(M_{\text{IL}}\)) across environments (Equation 8). Similar idea can be found in [17]. However, instead of training an invariant model given known environments, we freeze the invariant predictor and regularize the adjustment of \(\mathcal{E}_{\text{learn}}\) (i.e. the updates of \(M_{\text{EI}}\)) here.

\[L_{\text{IP}}=\text{Var}_{e\sim P(\mathcal{E}_{\text{learn}})}[\mathbb{E}^{e}[l (\Phi(X^{e}),Y^{e})]] \tag{8}\]

In summary, the training loss of \(M_{\text{EI}}\) can be summarized as _Environment Inference Loss_ (\(L_{\text{EI}}\)). The regularization strengths of \(L_{\text{LI}}\) and \(L_{\text{IP}}\) can be controlled by hyper-parameters \(\beta\) and \(\gamma\), respectively:

\[L_{\text{EI}}=L_{\text{ED}}+\beta L_{\text{LI}}+\gamma L_{\text{IP}} \tag{9}\]

In addition, before minimizing \(L_{\text{EI}}\), we pre-train our \(\Psi\) and one arbitrary \(f^{e}\) with ERM. In general, it empirically facilitates better feature extraction. Unlike EIIL [6] taking ERM as a reference model heavily relying on variant features, EDNIL performs more consistently under various choices of ERM. Namely, the initialization of EDNIL can be more arbitrary than that of EIIL. We verify the argument in Section 4.
### The Invariant Learning Model

To identify invariance across environments, IRM [1] is selected as our base algorithm optimizing the invariant predictor \(\Phi\) in our model \(M_{\text{IL}}\). As for the required environment partitions during training, we assign environment label \(e\in\operatorname{supp}(\mathcal{E}_{\text{learn}})\) with largest \(P(e|x_{i},y_{i})\), inferred by \(M_{\text{EI}}\) (Section 3.1.1), to each data \((x_{i},y_{i})\). 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 confidence score \(c_{e}=\mathbb{E}^{e}[P(e|X^{e},Y^{e})]\) for each environment \(e\in\operatorname{supp}(\mathcal{E}_{\text{learn}})\). Our training objective is modified to minimize _Invariant Learning Loss_ (\(L_{\text{IL}}\)) that considers the weighted average of environmental losses:

\[L_{\text{IL}}=\sum_{e\in\operatorname{supp}(\mathcal{E}_{\text{learn}})}w_{e }\cdot[R^{e}(\Phi)+\lambda||\nabla_{w|w=1.0}R^{e}(w\circ\Phi)||^{2}] \tag{10}\]

\[w_{e}=\frac{c_{e}}{\sum_{e^{\prime}\in\operatorname{supp}(\mathcal{E}_{ \text{learn}})}c_{e^{\prime}}} \tag{11}\]

## 4 Experiments

We empirically validate the proposed method on biased datasets, Adult-Confounded, CMNIST, Waterbirds and SNLI. The definition of spurious correlations mainly follows the protocols proposed by [1; 6; 9; 29]. In Section 4.1, Adult-Confounded and CMNIST are tested with Multilayer Perceptron (MLP). In Section 4.2, two more complex datasets, Waterbirds and SNLI, are used to evaluate the integration of transfer learning. Deep pre-trained models will be fine-tuned to extract representations.

The following four methods are selected as our competitors: Empirical Risk Minimization (ERM), Environment Inference for Invariant Learning (EIIL [6]), Kernelized Heterogeneous Risk Minimization (KerHRM [23]) and Invariant Risk Minimization (IRM [1], Equation 2). EIIL and KerHRM are invariant learning methods with unsupervised environment inference, which share the same settings as EDNIL. HRM [22] is replaced by KerHRM for non-linearity. For IRM who requires environment partitions, we re-label \(\mathcal{E}_{\text{oracle}}\) on each given biased training set, which diversifies the spurious relationships to elicit upper-bound performance of IRM.

For hyper-parameter tuning, we utilize an in-distribution validation set composed of 10% of training data. 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 as the evaluation metric.

Besides, more experimental details are revealed in Appendix B. We also discuss additional experiments in Appendix C, including solutions to regression problem and model stability given different spurious strengths at train time.

### Simple Datasets with MLP

This section includes two simple datasets, Adult-Confounded and CMNIST, where spurious correlations are synthetically produced with the predefined strengths. For all competitors, MLP is 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)2. Following [6], individuals are re-sampled according to sensitive features _race_ and _sex_ to simulate spurious correlations. Specifically, with binarized _race_ (Black/Non-Black) and _sex_ (Female/Male), four possible subgroups are constructed: Non-black Males (SG1), Non-black Females (SG2), Black Males (SG3), and Black Females (SG4). Keeping original train/test split and subgroup sizes from UCI Adult, we sample data based on the given target distributions in each sensitive subgroup as shown
in Table 2. 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 of 96 neurons is considered. For IRM, four environments comprise \(\mathcal{E}_{\text{oracle}}\), where the correlations between variant features \((race,sex)\) and target \(Y\) are distributed without overlapping. More details are provided in Appendix B.

ResultsThe results are shown in Table 3. 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. On the other hand, EIIL can improve consistency to some degree, but not as strong as EDNIL. One 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, as observed by Liu et al. [23].

Ablation Study for \(L_{\text{EI}}\)We first claim the importance of \(L_{\text{LI}}\), which constrains label dependency, by setting the coefficient \(\beta\) to zero. As discussed in Section 3, the resulting environments are determined purely by target labels, leading to inferior performance as shown in Table 3. Next, we demonstrate the effectiveness of joint optimization in Figure 2(a). The regularization \(L_{\text{IP}}\) promotes environment inference by considering existing invariant relationships, so that the worst-case performance improves and remains stable over iterations. According to Table 3, if the coefficient \(\gamma\) is turned off, feedback generated by \(M_{\text{IL}}\) will be ignored and the effect of invariant learning may become undesirable.

#### 4.1.2 Discussions on CMNIST

We report our evaluation on a noisy digit recognition dataset, CMNIST. Following [1], we first assign \(Y=1\) 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 labels \(C\), which are generated by randomly flipping \(Y\) with probability \(e\). For training, two equal-sized environments with \(e=0.1\) and \(e=0.2\) are merged, which is equivalent to one with \(e=0.15\) on average. For 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, we follow the settings recommended by [23]. Specifically, we randomly select 5,000 samples and train MLP with one hidden layer of 1024 neurons. To construct ideally diversified \(\mathcal{E}_{\text{oracle}}\) for IRM, we pack all examples with \(Y=C\) into one environments, and \(Y\neq C\) into the other.

ResultsThe results are shown in Table 4. First of all, not surprisingly ERM still adopts poorly to distributional shifts. Among all invariant learning methods without manual environment labels, EDNIL gets closest to IRM with \(\mathcal{E}_{\text{oracle}}\), achieving consistent and robust performance in this dataset. As shown in Figure 2(c), EDNIL provides almost ideally diversified \(\mathcal{E}_{\text{learn}}\) for invariant learning.

\begin{table}
\begin{tabular}{c c c c c} \hline \hline  & Train & \begin{tabular}{c} Test \\ (IID) \\ \end{tabular} & \begin{tabular}{c} Test \\ (IND) \\ \end{tabular} & 
\begin{tabular}{c} Test \\ (OOD) \\ \end{tabular} \\ \hline SG1 & 0.9 & 0.9 & 0.5 & 0.1 \\ SG2 & 0.1 & 0.1 & 0.5 & 0.9 \\ SG3 & 0.9 & 0.9 & 0.5 & 0.1 \\ SG4 & 0.1 & 0.1 & 0.5 & 0.9 \\ \hline \hline \end{tabular}
\end{table}
Table 2: \(P(Y=1|SG)\) for Adult-Confounded. IID shares spurious correlations with the train set. IND has no bias on _race_ and _sex_. OOD defines the worst-case performance.

\begin{table}
\begin{tabular}{c c c c} \hline \hline  & \multicolumn{1}{c}{IID} & IND & **OOD** \\ \hline ERM & \(\mathbf{92.4\pm 0.1}\) & \(\mathbf{66.8\pm 0.3}\) & \(\mathbf{40.7\pm 0.5}\) \\ \hline EIIL & \(\mathbf{76.2\pm 0.4}\) & \(\mathbf{73.5\pm 0.5}\) & \(\mathbf{70.2\pm 1.7}\) \\ KerHRM & \(\mathbf{82.4\pm 3.9}\) & \(\mathbf{75.1\pm 4.0}\) & \(\mathbf{67.9\pm 9.3}\) \\ KerHRM\({}^{*}\) & \(\mathbf{81.2\pm 1.8}\) & \(\mathbf{78.5\pm 0.3}\) & \(\mathbf{75.6\pm 1.9}\) \\ EDNIL & \(\mathbf{80.7\pm 0.4}\) & \(\mathbf{79.1\pm 0.4}\) & \(\mathbf{77.5\pm 0.3}\) \\ EDNIL\({}_{\beta=0}\) & \(\mathbf{91.8\pm 0.0}\) & \(\mathbf{66.7\pm 0.1}\) & \(\mathbf{41.3\pm 0.7}\) \\ EDNIL\({}_{\gamma=0}\) & \(\mathbf{78.2\pm 2.4}\) & \(\mathbf{75.4\pm 1.6}\) & \(\mathbf{72.5\pm 3.3}\) \\ \hline IRM & \(\mathbf{79.9\pm 0.4}\) & \(\mathbf{79.3\pm 0.3}\) & \(\mathbf{78.8\pm 0.4}\) \\ \hline \hline \end{tabular}
\end{table}
Table 3: Testing accuracy (%) on Adult-Confounded. Three subsets are defined in Table 2. \(\beta=0\) and \(\gamma=0\) indicate the removal of \(L_{\text{LI}}\) and \(L_{\text{IP}}\) when training \(M_{\text{EI}}\).
Number of EnvironmentsWhile the predefined number of environments is a hyperparameter that requires tuning, as illustrated in Figure 2(b), EDNIL can demonstrate a certain level of tolerance towards this parameter. Specifically, when the environment number is larger than the oracle (i.e. 2), some environment classifiers become redundant. Each of them provides a moderate constant loss, taking up fixed and ignorable space in the softmax function. The visualization of \(\mathcal{E}_{\text{learn}}\) with 5 available environments is shown in Figure 2(c). Additionally, training \(M_{\text{EI}}\) in EDNIL with more environments is much more efficient than clustering-based methods, such as the one proposed in KerHRM.

### 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 KerHRM due to its efficiency issue. In Section 4.2.1, image dataset, Waterbirds [29], with controlled spurious correlations is selected for evaluating the generalization on more high-dimensional images. In Section 4.2.2, a real-world NLP dataset, SNLI [3], is 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 on Waterbirds

In Waterbirds [29], each bird photograph, from CUB dataset [31], is combined with one background image, from Places dataset [33]. Both birds and backgrounds are either from land or water, and our target is to predict the binarized species of birds. At train time, landbirds and waterbirds are frequently present in land and water backgrounds, respectively. Therefore, empirically trained models are prone to learn context features, and fail to generalize as background varies [2; 6; 10; 21; 29].

To split an in-distribution validation set, we merge original training and validation data3 and split 10% for hyper-parameter tuning. For testing, we 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 target \(Y\) and background \(BG\), we distribute \(Y=BG\) and \(Y\neq BG\) into two different environments and apply balanced class weights for the oracle settings of IRM.

Footnote 3: In the original training split, backgrounds are unequally distributed in each class. However, they are equally distributed in the original validation split.

ResultsThe results are shown in Table 5. As observed in [6; 29], ERM suffers in the hardest case (i.e. waterbirds on land). E�IL also performs poorly in this case. With a more sophisticated learning framework, EDNIL narrows the gaps between subgroups and raises the worst-case performance. The results show that EDNIL can be more resistant to distributional shifts.

\begin{table}
\begin{tabular}{c c c c} \hline \hline Color Noise & 0.1 & 0.5 & **0.9** \\ \hline ERM & **88.4 \(\pm\) 0.3** & 55.0 \(\pm\) 0.5 & 21.7 \(\pm\) 0.8 \\ \hline E�IL & 79.6 \(\pm\) 0.3 & 71.7 \(\pm\) 0.7 & 63.1 \(\pm\) 0.5 \\ KerHRM & 74.3 \(\pm\) 0.7 & 66.2 \(\pm\) 1.7 & 58.0 \(\pm\) 11.5 \\ KerHRM* & 71.3 \(\pm\) 0.7 & 68.5 \(\pm\) 0.5 & 66.1 \(\pm\) 0.7 \\ EDNIL & 77.7 \(\pm\) 0.4 & **76.8 \(\pm\) 0.3** & **75.2 \(\pm\) 0.4** \\ \hline IRM & 77.8 \(\pm\) 0.4 & 76.8 \(\pm\) 0.4 & 75.2 \(\pm\) 0.3 \\ \hline \hline \end{tabular}
\end{table}
Table 4: Testing accuracy (%) of CMNIST, where color noise 0.9 makes the worst-case environment.

Figure 3: Analysis results. (a) Ablation study of joint optimization on Adult-Confounded, where \(\gamma=0\) means the removal of \(L_{\text{IP}}\). (b) CMNIST Performance of EDNIL with different configured numbers of environments. (c) Comparison between \(\mathcal{E}_{\text{learn}}\) inferred by EDNIL and the oracle environments \(\mathcal{E}_{\text{oracle}}\) on CMNIST. (d) Testing E�IL’s and EDNIL’s sensitivity to initialization on Waterbirds.
Choice of InitializationBoth EIIL and EDNIL take ERM as initialization. As mentioned in Section 1, heavy dependency on 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 in Figure 3d reveal the consistency of EDNIL, which accentuates our strength of less sensitive initialization. As implied by [6], EIIL could fail with a more well-fitted reference model in this case since ERM might get distracted from variant features. One is prone to be misled into an undesirable choice for EIIL when seeking hyper-parameters without prior knowledge of distributional shifts. For instance, the validation score of EIIL with a 500-step reference model (95.8%) slightly surpasses that of the 100-step model (94.7%), but their performance on testing exhibits significant discrepancy. On top of that, the selection of an effective reference model could be more intractable when spurious correlations in data are relatively weak. Supporting evidence can be observed in Appendix C.2.

#### 4.2.2 Discussions on SNLI

The target of SNLI [3] is to predict the relation between two given sentences, premise and hypothesis. Recent studies [12; 24; 27] reveal hypothesis bias in SNLI, which is characterized by patterns in hypothesis sentences highly correlated with a specific label. One can achieve low empirical risk without considering premises during prediction. However, as the bias no longer holds, the performance degradation occurs [11; 24].

We sample 100,000 examples and consider all classes, _entailment_, _neutral_ and _contradiction_, for our experiment. Following [9], we define three subsets by training a 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 [30] is taken as the pre-trained model for further mini-batch fine-tuning. For \(\mathcal{E}_{\text{oracle}}\), we assign the bias aligned subset to the first environment, and the misaligned one to the second. In order to make bias prevalence equal in the two environments, unbiased samples are scattered proportionally to the two environments.

ResultsThe results are shown in Table 6. As reported by [9], ERM receives higher score on the bias aligned subset, but it fails in 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 a high level, EDNIL is still capable of encoding and eliminating possible variant features.

\begin{table}
\begin{tabular}{c c c c} \hline \hline Subset & Unbiased & Aligned & **Misaligned** \\ \hline ERM & \(\mathbf{74.6\pm 0.3}\) & 94.7 \(\pm\) 0.2 & 52.6 \(\pm\) 0.9 \\ \hline EIIL & \(74.2\pm 0.3\) & \(\mathbf{95.0\pm 0.1}\) & 51.7 \(\pm\) 1.3 \\ EDNIL & \(74.3\pm 0.8\) & 94.2 \(\pm\) 0.2 & \(\mathbf{54.5\pm 1.0}\) \\ \hline IRM & \(74.0\pm 0.9\) & 92.3 \(\pm\) 0.5 & 56.9 \(\pm\) 1.1 \\ \hline \hline \end{tabular}
\end{table}
Table 6: Testing accuracy (%) on SNLI, where the misaligned defines the worst-case performance.

\begin{table}
\begin{tabular}{c c c c c} \hline \hline (Y, BG) & (Land, Land) & (Water, Water) & (Land, Water) & **(Water, Land)** \\ \hline ERM & 99.4 \(\pm\) 0.0 & \(\mathbf{91.4\pm 0.2}\) & \(\mathbf{90.9\pm 0.8}\) & 72.8 \(\pm\) 1.0 \\ \hline EIIL & \(\mathbf{99.4\pm 0.3}\) & 90.5 \(\pm\) 1.8 & 89.3 \(\pm\) 3.7 & 68.6 \(\pm\) 5.4 \\ EDNIL & 98.5 \(\pm\) 0.6 & 89.9 \(\pm\) 1.5 & 90.3 \(\pm\) 3.0 & \(\mathbf{78.6\pm 4.3}\) \\ \hline IRM & 98.0 \(\pm\) 0.5 & 90.6 \(\pm\) 1.1 & 89.5 \(\pm\) 1.7 & 83.2 \(\pm\) 2.2 \\ \hline \hline \end{tabular}
\end{table}
Table 5: Testing accuracy (%) of Waterbirds. The subgroup (Water, Land) contributes the worst-case performance.
Limitation

Our learning algorithm for environment inference is based on the graphical model plotted in Figure 2. However, it is important to acknowledge that data may not always conform to the assumed process, resulting in potential biases that cannot be adequately captured by the proposed neural network. In the paper, while we provide empirical studies of effectiveness on diverse datasets, we are still aware that a stronger guarantee of performance is required.

## 6 Conclusion and Societal Impacts

This work proposes EDNIL for training models invariant to distributional shifts. To infer environments without supervision, we propose a multi-head neural network structure to identify and diversify plausible environments. With joint optimization, the resulting invariant models are shown to be more robust than existing solutions on data having distinct 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, we 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 on debiasing machine learning systems. 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, as 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 a potential to learn a fair predictor without prior knowledge of sensitive attributes, which is related to [13; 19; 32]. We expect that one can extend our work to more fairness benchmarks and criteria in the future.

Last but not 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 a 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-E-002 -146 -MY3 and 110-2634-F-002-050 -.