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 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.
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 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 causal 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 Environments IRM 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 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 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.
As the distributional shifts are assumed to be unknown in the training stage, appropriate initialization is 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 it 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 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 does not have the three limitations discussed above.
The comparisons between EDNIL and other methods are shown in Table 1.
[Table]
[Equation]
(Ours) ✓ ✓ ✓
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, we discuss recent studies investigating ideal environments.
In Section 2.3, we introduce the existing unsupervised methods inferring environments.
2.1 Out-of-distribution Generalization and Invariant Learning
Following [1, 22], we consider a dataset D
= { D e } e ∈ supp( E tr ) with different sources D e = { ( x ei , y ei ) } n e i =1 collected under multiple training environments e ∈ supp( E tr ) .
Random variable E tr indicates the training environment labels.
For simplicity, X e , Y e and P e denote data, target label and distribution in environment e respectively.
With E all containing all possible environments such that supp( E all ) ⊃ supp( E tr ) , the goal of outof-distribution generalization is to learn a predictor f ( · ) :
X → Y as Equation 1, where R e ( f ) = E
X e ,Y e [ l ( f ( X e ) , Y e )] = E e [ l ( f ( X e ) , Y e )] measures the risk under environment e with any loss function l ( · , · ) .
In general, for e ∈ supp( E tr ) and e ′ ∈ supp( E all ) \ supp( E tr ) , P e ′ ( X, Y ) is rather different from P e ( X, Y ) .
[Equation]
Recently, several studies [1, 4, 18, 26, 28] attempt to tackle the generalization problems by discovering invariant relationships across all environments.
A commonly proposed assumption is the existence of invariant features X c and variant features X v .
Specifically, raw features X are assumed to be composed of X c and X v , or X
= h ( X c , X v ) where h ( · ) is a transformation function.
Invariant features X c are assumed to be equally informative for predicting targets Y across environments e ∈ supp( E all ) .
On the contrary, the distribution P e ( Y | X v ) can arbitrarily vary across e .
As a result, predictors depending on X v can have unpredictable performance in unseen environments.
In particular, the correlations between X v and Y are known as spurious and unreliable.
To extract X c , IRM [1] assumes there is an encoder Φ for obtaining representations Φ( X ) ≈ X c .
The encoder is trained with a regularization term enforcing simultaneous optimality of the predictor w ◦ Φ( X ) in training environments, where dummy classifier w = 1 . 0 is a fixed multiplier for the encoder outputs: (cid:88)
[Equation]
As w is a dummy layer, the encoder Φ is also regarded as a predictor.
2.2 Ideal Environments
As E tr is unavailable or sub-optimal in most applications, learning to find appropriate environments (denoted by E learn ) is attractive.
However, the challenge is the lack of knowledge of valid environments.
Recently, Lin et al.
[20] propose 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.
[Equation]
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 v ) , between environments is favored.
Given environments with a clear indication of distributional shifts, IRM can thus 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.
2.3 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 given 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. [Figure]
To extend HRM to a broader class of applications and improve the model performance, Liu et al.
[23] proposes 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.
In particular, 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.
Methodology
In this section, we propose a general framework to learn invariant model without manual environment labels.
As shown in Figure 1a, our proposed method consists of two models, M EI and M IL .
Given the pooled data ( X, Y ) , M EI infers environments E learn satisfying Condition 3 and 4, and M IL is an invariant model trained with the inferred environments.
Our framework is jointly optimized with alternating updates.
The learned M IL can provide information of invariant features to M EI , so that Condition 3 and 4 can be fulfilled simultaneously.
Note that M EI only serves at train time to provide environments for invariant learning.
At test time, only M IL is needed to perform invariant predictions.
3.1 The Environment Inference Model
The target of environment inference is to partition data into environments E 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 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 (Ψ( X )) , where Ψ( X ) is learned representations expected to encode variant features X 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 1b; particularly, a single-head network in M EI with shared parameters corresponds to a cluster center in HRM or KerHRM.
The training procedure of M EI can be divided into inference stage and learning stage .
3.1.1 Inference Stage of M EI
The goal is to infer an environment label for each training data.
As in the graphical model, E learn is associated with variant relationships.
Inspired by multi-class classification problem, we propose Equation 5, where the probability P ( e | X, Y ) is estimated via a softmax of negative l ( f e (Ψ( X )) , Y ) divided by a constant temperature τ .
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 (Ψ( X )) for each environment e . Intuitively, each data prefers the environment whose model has better prediction.
[Equation]
3.1.2 Learning Stage of M 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 ED ), Label Independency Loss ( L LI ) and Invariance Preserving Loss ( L IP ).
In particular, L ED and L IP correspond to the concepts of Condition 4 and 3 respectively.
Environment Diversification Loss ( L ED )
We consider maximizing H ( Y | X v ) − H
( Y | X v , E learn ) to satisfy Condition 4 and capture more diverse variant relationships.
Given the estimated P ( E learn | X, Y ) ,
L ED selects the most probable environment and its corresponding network for optimization:
[Equation]
For each data ( x i , y i ) , although only one environment e i = arg max e P ( e | x i , y i ) is selected for the minimization, the softmax simultaneously propagates gradient to maximize l ( f e ′ (Ψ( x i )) , y i ) for e [-0 
-]′ ̸
= e i .
The network learns to maximize the dependency between E 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 Independency Loss ( L LI )
With d-separation [25], E learn is independent of Y in the graphical model.
Hence, L LI constraints their dependency measured by the mutual information
I ( Y ; E learn ) .
Empirically, L LI also prevents a trivial solution that environments are determined purely by target labels Y regardless of input features X , which is undesirable for invariant learning.
To minimize I ( Y ; E learn ) , it can be verified that it is equivalent to minimizing Equation 7 given P ( Y ) is known.
[Equation]
Invariance Preserving Loss ( L IP ) For
Condition 3, as M IL learns some invariant relationships after several training steps (Section 3.2), L IP can be considered to exclude invariant features from the diversification.
Specifically, designed as the contrary of L ED , L IP limits the variance of expected loss from invariant predictor Φ (in M IL ) across environments (Equation 8).
Similar idea can be found in [17].
However, instead of training invariant model given known environments, we freeze the invariant predictor and regularize the adjustment of E learn (i.e. the updates of M EI ) here.
[Equation]
In summary, the training loss of M EI can be summarized as Environment Inference Loss ( L EI ).
The regularization strengths of L LI and L IP can be controlled by hyper-parameters β and γ respectively:
[Equation]
In addition, before minimizing L EI , we pre-train our Ψ and one arbitrary f e with ERM.
In general, it empirically facilitates better feature extraction.
Unlike [6] taking ERM as a reference model heavily relying on variant features, EDNIL performs more consistently under various choices of ERM.
In particular, it is not confined to under-fitted ERM training.
We verify the argument in Section 4.
3.2 The Invariant Learning Model
To identify invariance across environments, IRM [1] is selected as our base algorithm optimizing the invariant predictor Φ in our model M IL .
As for the required environment partitions during training, we assign environment label e ∈ supp( E learn ) with largest P ( e | x i , y i ) , inferred by M 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 the confidence score c e for each environment e ∈ supp( E learn ) , i.e. E e [ P ( e | X e , Y e )] .
Our training objective is modified to minimize Invariant Learning Loss ( L IL ) that considers the weighted average of environmental losses:
[Equation]
Experiments
We empirically validate the proposed method on biased datasets, Adult-Confounded, CMNIST,
Waterbirds and SNLI.
The generation of spurious correlations mainly follow 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 taken for evaluating the integration of representation learning.
Modern deep models with pretrained weights will serve as the encoders of variant and invariant features.
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 E oracle on each given biased training set, which diversifies the spurious relationships to elicit upper-bound performance of IRM.
For hyper-parameter 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 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.
4.1 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 choice of income threshold. [Table]
[Table] [Table] 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 E oracle , where the correlations between variant features ( race, sex ) and target Y are distributed without overlapping.
More details are provided in Appendix B.
Results The 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.
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, as observed by Liu et al.
Ablation Study for L EI
We first claim the importance of L LI , which constraints label dependency, by setting the coefficient β to 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 3.
Next, we demonstrate the effectiveness of joint optimization in Figure 3a.
The regularization L IP promotes environment inference, so that the worst-case performance improves and remains stable over iterations.
[- = 
-]
According to Table 3, if the coefficient γ is turned off, feedback generated by M IL will be ignored and the effect of invariant learning risks being 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 label C , which is 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
E oracle for IRM, we pack all examples with C =
Y into one environments, and C [-
-]̸
= Y into the other.
Results The 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,
[Figure] EDNIL gets closest to IRM with
As shown in Figure 3c EDNIL provides almost ideally diversified E learn for invariant learning.
[Table] Number of Environments As shown in Figure 3b, 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 redundant.
Each of them provides a moderate constant loss, taking up fixed and ignorable space in the softmax function.
The visualization of E learn containing 5 available environments is shown in Figure 3c.
Additionally, training M EI in EDNIL with more environments is much more efficient than clustering-based methods, such as the one proposed in KerHRM.
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 [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.
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 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 a validation set whose overall distribution is i.i.d.
to the training set, we merge original training and validation data 3 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 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.
Results The results are shown in Table 5.
As observed in [6, 29], ERM suffers in the hardest case (i.e. waterbirds on land).
EIIL also performs poorly in this case.
With a more sophisticated learning [Table] 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 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 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 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 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.
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, unbiased, bias aligned and bias misaligned, 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.
[Table]
For all methods, DistilBERT [30] is taken as the pre-trained model for further mini-batch finetuning.
For E oracle , we assign the bias aligned subset to the first environment, and the bias misaligned subset to the second.
In order to make bias prevalence equal in the two environments, unbiased samples are scattered proportionally to the two environments.
Results The 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 diversifying possible variant features.
Limitations
Our learning algorithm for environment inference is based on the graphical model plotted in Figure 2.
As data are not necessarily generated by the presumed process, there can exist biases that cannot be captured by the proposed neural network.
In the paper, we provide empirical studies of effectiveness on diverse datasets, while we are still aware that a stronger guarantee of performance is required.
Conclusions 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.
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, 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. Appendix A Proof of the Underlying Graphical Model of EDNIL We assume data are generated by Equation 12 which is the process adopted by Arjovsky et al.
[1] and
Lin et al.
[20], where g c , g v are deterministic functions, and ϵ c , ϵ v are random noises independent of X c and X v .
Y =
g c ( X c , ϵ c ) ,
X v = g v ( Y, ϵ v ) . Equation 12 establishes a chain that X c → Y → X v . With the additional relation E → X v proposed in our graphical model, the following two dependencies can be obtained via d-separation [25]:
[Equation]
As can be seen, the conditional independence between Y and E given X c leads to H ( Y | X c ) = H ( Y | X c , E ) .
On the other hand, the dependency between Y and E given X v implies that there exists an environment variable satisfying H ( Y | X v ) > H ( Y | X v , E ) . B Experimental Details B.1 Computing Resources
All experiments were run on a GeForce RTX 3090 machine. The training time and GPU memory consumption of EDNIL are specified in Table 7. It takes approximately 30 hours for EDNIL to accomplish all tasks, including main experiments and analyses. For the choices of deep encoders, we utilize Resnet34 from torchvision 4 on Waterbirds, and
DistilBERT from Huggingface 5 on SNLI. Both network architectures and pre-trained weights are kept as the default. [Table] [Table]
[Table]
[Table]
Without leaking into out-of-distribution information, 10% of training data are split as an in-distribution validation set.
Particularly, we infer environments on the validation data with M EI and determine hyper-parameters according to the worst-environment score.
For M EI , we select number of environments from 2 to 5, τ in softmax function from 0.05 to 0.5, β and γ in L EI from 0.2 to 10.
In L ED , we clip overaggressive w i with an upper bound w thres for training stability, and it is chosen from 1.2 to 5.
As for M IL , the penalty strength λ in L IL is selected from {2, 10, 100, 1000}.
Following Arjovsky et al.
[1], we conduct an annealing mechanism before using the configured penalty strength.
The chosen number of annealing iterations ranges from 20% to 80% of the whole.
For the complex datasets, we consider a longer annealing period (larger than 50%) to learn basic representations better.
Number of total training steps is decided from 500 to 2000.
In Adult-Confounded and CMNIST, full-batch training is implemented due to enough memory space.
In Waterbirds and SNLI, batch size is chosen among 128, 256 and 512.
The choices of learning rate and optimizer depend on the dataset.
For Adult-Confounded, CMNIST and Waterbirds, a learning rate between 2e-4 and 2eis considered.
We take Adam as the optimizer for Adult-Confounded and CMNIST, and choose
SGD for Waterbirds.
For SNLI, a smaller learning rate in {2e-5, 3e-5, 5e-5, 1e-4} is selected when fine-tuning DistilBERT with AdamW.
B.3 Oracle Settings on Adult-Confounded
Given a biased training set and two sensitive features, race and sex , E oracle for IRM is constructed according to Table 8.
The correlations between variant features and target are maximized within each environment and diversified across environments.
As implied by [5, 22], spurious correlations are supposed to be eliminated when an invariant learning algorithm converges properly.
where e i ∈ supp( E oracle ) represents the i -th environment.
[Table]
B.4 Biased Model for SNLI
To define subsets for evaluation on SNLI, we follow the labeling procedure in [9].
Specifically, k-fold cross validation ( k = 5 ) is applied on the training set.
We fine-tune BERT [7] with hypothesis as its only inputs on k − 1 folds, and score the left-out k -th set.
For the development and testing sets, we score each example with average predictions from k different models.
The accuracy of the biased model is approximately 68%.
Finally, we set two thresholds ( t 1 , t 2 ) , defined by [9], to (0.2, 0.5), where t 1 is used to identify unbiased data, and t 2 is used to define bias aligned and misaligned sets. C Additional Empirical Results C.1 Synthetic Data for Regression Problem We further validate our work on regression problem with a synthetic dataset proposed in [23]. The features are X
= H [ X c , X v ] ∈ R d , and the target is generated by Y = f ( X c ) + ϵ , where H is an random orthogonal matrix for scrambling features and f ( · ) is a non-linear function.
X c are invariant features
that P ( Y | X c ) is consistent across environments, while X v are variant features that P ( Y | X v ) can arbitrary change according to the following data sampling mechanism:
[Equation]
In Equation 14 where | r | > 1 , r controls the spurious correlation between the certain variable
X ∗ v ∈ X v and the target label Y . Specifically, larger | r | represents stronger correlation, and the sign of r indicates the direction of correlation. In the training set, there are 1000 examples generated from the environment with r = 2 . 3 and 100 examples from that with r = − 1 . 1 . The environment labels are unavailable as in the previous experiments.
For testing, two scenarios are considered. First, we define two environments, IID and OOD, to evaluate the generalization under dramatic distributional shifts. In IID, 1100 examples are sampled with the same procedures as training data.
In OOD,examples are generated from the environment with r = − 2 . Secondly, following [23], we evaluate the stability over six testing environments where r ∈ {− 2 .
9 , − 2 . 7 , − 2 . 5 , ..., − 1 . 9 } .
In our regression task, the evaluation metric is mean square error.
For all methods, MLP with one hidden layer of 1024 neurons is utilized.
When calculating the label independency term L LI for
EDNIL, we discretize Y by quartiles. Empirically, such efficient estimation can improve the quality of environment inference to some degree.
We leave more precise approximations of mutual information between discrete and continuous variables for future work.
Results
The results of the first testing scenario are listed in Table 9. Among all methods, EDNIL obtains the most consistent scores across IID and OOD. The performance degradation when β =suggests the importance of L LI . Table 10 shows the results of the second scenario.
As in [23], Mean Error and Std Error represent the mean and standard deviation of errors over six testing environments respectively, both of which are averaged over 20 runs. Similar to the first scenario, EDNIL performs the best and most robustly in out-of-distribution settings. The estimated L LI also gains empirical improvements in this test.
[Table]
[Table] C.2 CMNIST with Different Color Noises
In this task
[Table]
the learning effects of all methods are tested under different strengths of spurious correlation at train time. We select CMNIST with fixed label noise 0.2, and adjust overall color noise e from 0.1 to 0.3 to generate five different training sets. For testing, e = 0 . 1 and e = 0 .
9 are considered. Given color label C and target Y , samples with C = Y are more than those with C ̸
= Y when e = 0 .
1 , where the direction of spurious correlation is aligned with that in the training sets. On the other hand, when e = 0 . 9 , samples with C ̸ = Y are in the majority.
Due to the reversed correlation, models relying on the variant feature (i.e. color C ) are vulnerable to this setting.
Results The results are plotted in Figure 4. As the training color noise increases, the generalization of ERM improves since the spurious correlation decreases at train time. Meanwhile, as indicated in [6], EIIL fails because the reference model, i.e. ERM, is no longer a pure variant predictor.
For
KerHRM, the instability of inferred environments results in large standard deviation, especially when training color noise is small.
In comparison, given different strengths of spurious correlation for training, EDNIL can distinguish invariant features from variant ones, and performs more consistently as IRM (with E oracle ) does.
[Figure]