Environment Diversification with Multi-head Neural Network for Invariant Learning

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.

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 ($\text{{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:

For each data $(x_{i},y_{i})$, although only one environment $e_{i}=\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 ($\text{{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.

Invariance Preserving Loss ($\text{{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.

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:

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.

We take UCI Adult [16] to predict binarized income levels (above or below $50,000 per year)²²2UCI Adult dataset is widely used in algorithmic fairness papers. However, a recent study [8] discusses some limitations of this dataset, such as the choice of income threshold.. 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 3. 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.

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. 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 $\text{{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 3(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.

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.

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, EDNIL gets closest to IRM with $\mathcal{E}_{\text{oracle}}$, achieving consistent and robust performance in this dataset. As shown in Figure 3(c), EDNIL provides almost ideally diversified $\mathcal{E}_{\text{learn}}$ for invariant learning.

Number of Environments While the predefined number of environments is a hyper-parameter that requires tuning, as illustrated in Figure 3(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 3(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.

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 data³³3In the original training split, backgrounds are unequally distributed in each class. However, they are equally distributed in the original validation split. 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.

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 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.

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 in Figure 3(d) 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.

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.

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 eliminating possible variant features.