Learning an Invertible Output Mapping Can Mitigate Simplicity Bias in Neural Networks
Abstract
Deep Neural Networks (DNNs) are known to be brittle to even minor distribution shifts compared to the training distribution.
While one line of work has demonstrated that Simplicity Bias (SB) of DNNs – bias towards learning only the simplest features – is a key reason for this brittleness, another recent line of work has surprisingly found that diverse/ complex features are indeed learned by the backbone, and their brittleness is due to the linear classiﬁcation head relying primarily on the simplest features.
To bridge the gap between these two lines of work, we ﬁrst hypothesize and verify that while SB may not altogether preclude learning complex features, it ampliﬁes simpler features over complex ones.
Namely, simple features are replicated several times in the learned representations while complex features might not be replicated.
This phenomenon, we term Feature Replication Hypothesis , coupled with the Implicit Bias of SGD to converge to maximum margin solutions in the feature space, leads the models to rely mostly on the simple features for classiﬁcation.
To mitigate this bias, we propose Feature Reconstruction Regularizer (FRR) to ensure that the learned features can be reconstructed back from the logits.
The use of FRR in linear layer training ( FRR-L ) encourages the use of more diverse features for classiﬁcation.
We further propose to ﬁnetune the full network by freezing the weights of the linear layer trained using FRR-L , to reﬁne the learned features, making them more suitable for classiﬁcation.
Using the proposed approach, we demonstrate noteworthy gains on synthetic/ semi-synthetic datasets, and outperform existing SOTA on the standard OOD benchmark DomainBed as well.
Introduction
DNNs are brittle against even minor shifts in the data distribution during inference, which are not uncommon in a real world setting [1, 2].
In this work, we aim to tackle the problem of Out-OfDistribution (OOD) generalization of Neural Networks in a covariate-shift [3] based classiﬁcation setting, by addressing the fundamental cause of their brittleness, rather than by explicitly enforcing invariances in the network using domain labels or data augmentations.
More speciﬁcally, we aim to mitigate the issue of Simplicity Bias , which is the tendency of Stochastic Gradient Descent (SGD) based solutions to overly rely on simple features alone, rather than on a diverse set of features [4, 5].
While this behavior was earlier used to explain the remarkable generalization of Deep Networks, recent works suggest that this is indeed a key reason behind their brittleness to domain shifts [6].
The extent of Simplicity Bias seen in models is a result of two important factors - diversity of features learned by the feature extractor, and the extent to which these diverse features are used for the task at hand, such as classiﬁcation.
Recent works suggest that generalization to distribution shifts can be improved by retraining the last layer alone, indicating that the features learned may already be good enough for the same [7, 8].
Does this imply that brittleness of models can be attributed to the learning of the classiﬁcation head alone?
If this is the case, why does SGD fail to utilize these diverse features despite its Implicit Bias to converge to a maximum margin solution in a linearly separable case [9]?
To answer these questions, we ﬁrstly hypothesize and empirically verify that Simplicity Bias leads to the learning of simple features over and over again, as compared to other, more complex features.
For example, among the 512 penultimate layer features of a ResNet, 462 of them might capture a simple feature such as color , while the remaining 50 might capture a more complex feature such as shape – we refer to this as (Simple) Feature Replication Hypothesis .
Assuming feature replication hypothesis, we further show theoretically and empirically that a maximum margin classiﬁer in the replicated feature space would give much higher importance to the replicated feature when compared to others, highlighting why the linear layer relies more on simpler features for classiﬁcation.
To mitigate this, we propose a novel regularizer termed Feature Reconstruction Regularizer (FRR) , to enforce that the features learned by the network can be reconstructed back from the logit or presoftmax layer used for the classiﬁcation task.
As shown in Fig.1, we ﬁrstly propose to train the linear classiﬁer alone by freezing the weights of the feature extractor.
This formulation enables the learning of an Invertible Mapping in the output layer, speciﬁcally for the domain of features seen during training.
This further allows the logit layer to act as an information bottleneck, encouraging all the factors of variation in the features to be utilized for the classiﬁcation task, thereby improving the diversity of features used.
We theoretically show that adding this constraint while ﬁnetuning the linear layer can learn a max-margin classiﬁer in the original input space, disregarding feature replication.
Consequently, the learnt linear classiﬁer also gives more importance to non replicated complex features while making predictions.
We further explore the possibility of improving the quality of features learned by the feature extractor, by using FRR for ﬁnetuning the backbone as well.
In order to do this, we freeze the linear classiﬁcation head, and further ﬁnetune the backbone with
FRR.
We ﬁnd that this encourages the network to indeed learn better quality features that are more relevant for classiﬁcation.
We list the key contributions of this work below -
• Key Observation : We provide a crisp hypothesis of “feature replication” to explain the brittleness of ERM trained neural networks to OOD data (Sec 2).
Using this, we provide theoretical and empirical evidence to justify the existence of Simplicity Bias in max margin classiﬁers (Sec 2.2). • Novel Algorithm based on the Observation : Based on this, we introduce a novel FRR regularizer to safeguard against the feature replication phenomenon.
Furthermore, we introduce a simple FRRL method to only regularize the linear head with FRR, and then introduce FRR-FLFT training regimen to train the feature extractor for improved OOD robustness (Sec 2.1).
We also provide theoretical support for FRR in an intuitive data distribution setting (Sec 2.2).
• Empirical validation of the hypothesis and the proposed algorithm : We demonstrate the effectiveness of FRR-FLFT and FRR-L by conducting extensive experiments on semi-real datasets (Table 3) constructed to study OOD brittleness, as well as on standard OOD generalization benchmarks, where FRR-FLFT can provide up to 3% gains over SOTA methods (Table 1).
Feature Replication Hypothesis and Feature Reconstruction Regularizer
Prior works have shown that neural networks trained with SGD exhibit simplicity bias (SB), even when initialized with pre-trained models that can capture complex features.
Our Feature Replication
Hypothesis – FRH– states that: SB is observed because the simpler features of the input are replicated multiple times in the feature space of neural networks.
When trained using SGD, the ﬁnal linear layer then learns the max margin classiﬁer on these replicated features, which leads to over-reliance on simpler features in the input.
Hence, the outputs of the network are brittle to distribution shifts that change such replicated features.
In this section, we propose a new regularizer – FRR– to mitigate this effect and further provide empirical and theoretical evidence for FRH and FRR
2.1 Method: Feature Reconstruction Regularizer (FRR)
To alleviate the simple feature replication issue, we propose Feature Reconstruction Regularizer (FRR) to enforce that the learned features can be reconstructed from the output logits.
We retrain the ﬁnal linear layer using this regularizer so that the model can utilize diverse features to compute the ﬁnal output.
To implement this, we introduce another neural network with the objective of reconstructing the features from the output logits, i.e. features f θ ( x ) should be recoverable from the predictions of the network through a transform T φ ( . ) parameterized by φ as shown below:
[Equation]
[Figure]
(a)
SVM, 5-Rep   (b)
FRR (Ours), 5-Rep where || .
||
p denotes the (cid:96) p norm.
We set this norm to be (cid:96) ∞ or (cid:96) 1 in our experiments.
Note that in order to ﬁnd the appropriate φ , we jointly optimize W the weight matrix of the linear classiﬁer, and φ using gradient descent based optimizers.
Our training pipeline consists of three stages, as summarized in Fig 1 and Alg 1.
We pretrain our model using the standard cross entropy loss L std , which is followed by retraining only the ﬁnal linear (and reconstruction) layer using our proposed regularizer L FRR .
Finally, we ﬁx the linear layer and ﬁnetune the backbone using L FRR [10], to learn features which are more useful for classiﬁcation.
2. FRH & FRR: Empirical and Theoretical Analysis
Empirical Justiﬁcation: In order to empirically demonstrate feature replication, we consider a binarized version of the coloured MNIST dataset [11] in Appendix B.2.
We observe that ERM learns more colour (simple) features than shape (complex) features, and the prediction is less correlated with the shape features.
We further empirically validate FRR on Coloured MNIST, where using FRR with the linear layer leads to lower correlation with Colour and higher correlation with shape when compared to standard ERM (Table 2).
Consequently, OOD accuracy improves by 5% over ERM.
Theoretical Analysis: We now present a simple data distribution with feature replication that highlights the OOD brittleness of standard ERM, and also demonstrates the effectiveness of FRR.
Data Distribution : Consider a linearly separable distribution consisting of two factors of variation as shown in Figure 2.
That is, consider the following distribution ( x, y ) ∼ D , where, y = ± 1 with probability 0 .
[Equation]
Also consider a feature extractor f θ ( . ) which captures feature replication in the ﬁrst feature, i.e. for every data point ( x, y ) , the new, feature replicated data point will be (˜ x, y ) , where,
[Equation]
i.e., x 1 is repeated d times.
The joint distribution of features and labels is denoted by ˜ D .
Finally, we deﬁne the l 2 max margin classiﬁer over a distribution D as w MM :=
arg min w 12 (cid:107) w (cid:107) 22 subject to y · (cid:104) w, x (cid:105) ≥ 1 ∀ ( x, y
) ∈
Supp ( D ) .
Then we have the following results:
Claim 2.1 (Brittleness due to Feature Replication) .
Consider the data distribution given in Equation 2, 3.
Then, the following holds: (1.)
The max-margin classiﬁer w MM over D is given by w MM = [1 , 1] , and (2.)
The max-margin classiﬁer ˜ w MM over ˜ D is given by ˜ w MM = (cid:2) 2 d +1 , · · · , 2 d +1 (cid:3) ∈ R d +1 .
The above claim implies that when there are replicated features to the input of the linear layer, the max-margin classiﬁer would give much more importance to the feature that is replicated.
Hence, even a minor change in this replicated feature in the input space would be ampliﬁed in the output of the classiﬁer.
This is especially concerning in light of the observations in Table 2, which validate the
Feature Replication Hypothesis in Coloured MNIST.
Proposition 2.2 (Robustness of FRR) .
Denote the average feature reconstruction loss L FRR ( ˜ w, ˜ φ ) := max 1 ≤ i ≤ d +1 E (˜ x,y ) ∼ ˜ D
(cid:104)
( (cid:104) ˜ w, ˜ x (cid:105) ˜
φ i − ˜ x i ) 2 (cid:105) and consider any ( ˜ w ∗ , ˜ φ ∗ ) satisfying:
[Equation]
[Equation]
arg min ( ˜ w, ˜ φ ) L FRR ( ˜ w, ˜ φ ) subject to y · (cid:104) ˜ w, ˜ x (cid:105) ≥ 0 ∀ (˜ x, y ) ∈ Supp (cid:16) ˜ D (cid:17) .
We have that: ˜ w ∗ 1 + · · · + ˜ w ∗ d = ˜ w ∗ d +1 .
Consequently, we have (cid:104) ˜
w ∗ , ˜ x (cid:105) ∝ (cid:104)
w MM , x (cid:105) for all x ∈ R 2 .
[Table]
The above result shows that the feature reconstruction regularizer will produce a linear classiﬁer that gives equal weights to the replicated and non-replicated features.
This is equivalent to a maximum margin classiﬁer in the non-replicated feature space , thereby resulting in enhanced robustness to distribution shifts.
Same is reﬂected in Figure 2 (b) which shows the impact of FRR on the trained boundary in the non-replicated feature space (See Appendix B.1).
We defer the proofs to Appendix A.
Experimental Results
We demonstrate the effectiveness of the proposed Feature Reconstruction Regularizer towards mitigating Simplicity Bias by evaluating the same on a 10-class variant of the semi-synthetic MNISTCIFAR dataset [6] in Appendix D. We further show the efﬁcacy of FRR towards improving OOD generalization on the DomainBed [11] benchmark.
We use the performance of the model on indomain validation data (i.e. the in-domain strategy by [11]) to select the best hyper-parameters, and report the average performance and standard deviation across 5 random seeds.
We compare our method against standard ERM training, which has proven to be a frustratingly difﬁcult baseline [11], and also against several state of the art methods on this benchmark - SWAD [12], MIRO [13] and
SMA [14] (See Appendix F).
Finally, we show that our approach can be effectively integrated with stochastic weight averaging to obtain further gains.
See Appendix G for further experimental details.
Main Results: The main results of our algorithm are reported in Table 1.
We ﬁnd that our pipeline of training and ﬁnetuning with FRR, when combined with ERM achieves improved performance with respect to the state of the art methods that do not use model weight-averaging, and in fact achieves comparable performance to methods that use model weight averaging as well.
Further, our method obtains substantial gains of more than 3% over ERM across datasets.
This empirically validates our ﬁnetuning paradigm which we denote as ERM+FRR.
Finally, using our method in tandem with SWAD helps us achieve a new state-of-the-art on the DomainBed benchmark , better than existing SOTA by close to 1% on average.
We report detailed results and further ablations in Appendices H and I.
Conclusion and Discussion
In this work, we consider the problem of OOD generalization through the lens of mitigating Simplicity
Bias in Neural Network training.
To unravel the paradox pertaining to the existence of Simplicity
Bias in learning only the simplest features, and the observation that the features learned by large practical models may already be sufﬁciently diverse, we put forth the Feature Replication Hypothesis that conjectures the learning of replicated simple features and sparse complex ones.
Combining this with the Implicit Bias of SGD to converge to maximum margin solutions, we provide a theoretical justiﬁcation to the high OOD sensitivity of Neural Networks.
To speciﬁcally overcome the effect of simple feature replication in linear layer training, we propose the Feature Reconstruction Regularizer , that penalizes the (cid:96) p norm distance between the features and their reconstruction from the output logits, thus improving the diversity of features used for classiﬁcation.
We further propose to freeze the weights of the linear layer thus trained, and use the FRR regularizer for ﬁnetuning the full network, to reﬁne the features to be more useful for the downstream task.
We justify the proposed regularizer both theoretically and empirically on synthetic and semi-synthetic datasets, and demonstrate its effectiveness in a real world OOD generalization setting.
We believe and hope that this work can pave the way towards obtaining a better understanding on the underlying causes for OOD brittleness of neural networks, and inspire the development of better algorithms for addressing the same.
A Proofs of theoretical results
Proof of claim 2.
Proof.
First, observe that w =
[1 , 1] satisﬁes the constraint y ·(cid:104) w, x (cid:105) ≥ 1 ∀ ( x, y
) ∈
Supp ( D ) .
Now, consider any w satisfying this constraint.
By considering ( y · (0 . 5 , 0 . 5) , y ) ∈ Supp ( D ) , we have that w 1 + w ≥ 1 .
We now conclude by observing that among w which satisfy w 1 + w 2≥ 1 , [1 , 1] minimizes 12 (cid:107) w (cid:107) 22 .
The proof of the second part of the claim is identical to that of the ﬁrst part.
First, observe that ˜ w = (cid:2) {+
+} d +1 , · · · , 2 d +1 (cid:3) satisﬁes the constraint y ·(cid:104) ˜ w, ˜ x (cid:105) ≥ 1 ∀ (˜ x, y
) ∈ Supp (cid:16) ˜ D (cid:17) .
Now, consider any ˜ w satisfying this constraint.
By considering ( y · (0 . 5 , 0 . 5 , · · · , 0 . 5) , y ) ∈ Supp (cid:16) ˜ D (cid:17) , we have that ˜ w 1 + ··· + ˜ w d
+1≥
We now conclude by observing that among ˜ w which satisfy ˜ w 1 + ··· + ˜ w d +1≥ 1 , (cid:2)d +1 , · · · , 2 d +1 (cid:3) minimizes 12 (cid:107) ˜ w (cid:107) 22 .
The classiﬁer ˜ w MM ∈ R
d +1 on
˜ x ∈ R d +1 is equivalent to the classiﬁer w proj = (cid:2)
2 d d +1 , 2 d +1 (cid:3)[-∈ R 
-]
∈ R 2 on x ∈ R 2 : (cid:104) w proj , x (cid:105) =
(cid:104)
˜ w MM , ˜ x (cid:105) ∀ x ∈ R 2 .
Proof of Proposition [-2.
-]2.
Proof.
Given any ˜ w , we ﬁrst compute L FRR ( ˜ w ) :
= min ˜
φ L FRR ( ˜ w, ˜ φ ) .
To do so, we note that the minimizing ˜ φ for a given ˜ w , denoted by ˜ φ ∗ ( ˜ w ) is given by:
[Equation]
So, ˜ φ ∗ i ( ˜ w ) =
E [ (cid:104) ˜ w, ˜ x (cid:105) ˜
x
i ] E [ (cid:104) ˜ w, ˜ x (cid:105) 2 ] .
So, we have that:
[Equation]
Since (cid:104) ˜ w, ˜ x (cid:105) =
˜ w 1 → d · ( y + n 1 ) + ˜ w d +1 · ( y + n 2 ) , where ˜ w 1 → d := (cid:80) di =1 ˜ w i , we have:
[Equation]
if
[Equation]
˜
[Equation]
if
[Equation]
1 , where we used the fact that E [ n 21 ] =
E [ n 22 ] = 1 / 12 .
Consequently, we have that L FRR ( ˜ w ) =
1312 (cid:16) 1 − min( ˜ w 1 → d , ˜ w d +1 ) 2 ( ˜ w 1 → d + ˜ w d +1 ) 2 + ˜ w 21 → d / 12+ ˜ w 2 d +1 / 12 (cid:17) .
Consequently, any minimizer ˜ w ∗ of L FRR ( ˜ w ) satisﬁes (cid:80)
[Table]
B Synthetic datasets B.1 Toy dataset
In line with the theoretical formulation described in Section-2.2, we further empirically validate the brittleness of SVM models and the highlight the effectiveness of the proposed Feature Reconstruction
Regularizer in the presence of replicated features.
We consider a linearly separable toy distribution consisting of two factors of variation as shown in Figure 2.
We deﬁne the means of the two classes at (1,1) and (-1,-1) and construct 500 data points in each class by adding noise sampled from
Unif[-0.5,0.5] independently along each dimension to the respective means.
We sample an Outof-Distribution (OOD) test set from a Uniform distribution with means centered at (1,0) and (-1,0) respectively, and similar noise along each dimension as the train set.
Therefore, while the train distribution can be classiﬁed by considering either the features aligned with the X-coordinate or the
Y-coordinate, the test set performance crucially depends on the variation along X-coordinate alone.
We consider feature replication along the y-axis, and hence construct this OOD dataset to verify the extent to which the other feature is considered for classiﬁcation.
To select the best hyperparameter
[Table]
for both SVM and FRR, we consider the presence of a validation set whose distribution is similar to the test distribution.
As shown in Figure 2, we observe that the SVM model starts relying more on the replicated features alone in case of feature replication, compromising its performance on the
OOD data.
The proposed regularizer on the other hand, gives equal importance to both features even in the presence of feature replication, resulting in improved OOD generalization.
B.2 Empirical validation of Feature Replication Hypothesis (FRH) in ERM
Coloured MNIST dataset.
In order to empirically demonstrate feature replication, we use a binarized version of the coloured MNIST dataset [11].
To construct this dataset, we ﬁrstly assign two digits of the MNIST dataset, namely “1” and “5”, to classes 0 and 1 respectively.
For the in-domain training distribution, we associate colours in the range R 0 =
[(115 , 0 , 0) − (256 , 141 , 0))]
(i.e. red) to label
(i.e. the digit “1”) and the range R 1 =
[(0 , 115 , 0) − (141 , 256 , 0)]
(i.e. green) to the label(i.e.
the digit “5”), where colors are represented in the RGB space.
To summarize, while training the network, we super-impose images of “1” onto colours of range R 0 , and images of “5” onto colours of range R 1 .
It is to be noted that the choice of colour ranges as deﬁned above introduces an overlapping range between [(115 , 115 , 0) − (141 , 141 , 0))] where images are associated with labels 0 and 1 with equal probability.
This overlap reduces the correlation of colour features with labels, while shape features have a correlation of 1 with the labels.
In Figure 3, we show examples of images from the train and test distributions of this dataset.
In Figure 4, we pictorially depict the correlations between the 32 features learnt by the network.
We can see a block structure emerging, indicating that there is a high amount of feature replication.
Training setup : We train a model on this dataset, and test it on images which do not have any correlation between the label and the colour, i.e. images where the digits “1” and “5” are superimposed on randomly coloured backgrounds.
We construct this test distribution to see how well different algorithms learn simple (i.e. colour) and complex (i.e. shape) features, since an algorithm which depends only on the spurious colour features would not have good performance on the test domain.
We train a four layered CNN on this data.
If a feature in the penultimate layer f θ ( x ) has more than 90% correlation with the color or shape of the input, then we call it as a color feature or a shape feature, respectively.
We also compute the correlation of these features with the output of the network
( W T f θ ( x ) ) over inputs from the test domain.
This gives us information of the learnt features, and their contributions to the ﬁnal prediction of the network.
Note that the feature dimension is m = 32 , and the output dimension is k = 1 .
Observations :
In Table 2, we report the number of colour features, shape features, and the average correlation of each of these with the ﬁnal prediction.
We observe that the ERM trained model learns both shape and colour features, but the number of learnt colour features ( 26 ) is much higher than the number of shape features ( 4 ), despite their weaker correlation with labels, thus validating our Feature
Replication Hypothesis .
We also visualize the inter-feature correlation of the learnt features in Fig 4, which shows blocks of highly correlated features, further validating our hypothesis.
Furthermore, we note that correlation of the output with the shape features is lower, leading to OOD accuracy of 59% only.
C Algorithm
Algorithm 1: Our training algorithm
Data: Training data D S
= { ( x i , y i ) : i ∈ [ n ] } , model ( θ, W ) , feature reconstruction model φ , λ FRR , λ FT 1 θ std , W std ← Adam (min θ,W (cid:80)
i
L std ( θ, W, ( x i , y i ))) .
/* Standard training of model parameters θ and W .
*/ 2 Freeze θ to be θ std /* Initializing model for training with FRR.
*/ 3 W FRR , φ FRR ← Adam (min W,φ (cid:80) i
L std ( θ std , W, ( x i , y i ))
+ λ L L FRR ( x i , θ std , W, φ )) .
/* FRR-L: Training W, φ with FRR defined in eqn.
1 */ 4 θ FLFT ← Adam (min θ (cid:80) i
L std ( θ, W FRR , ( x i , y i ))
+ λ FLFT L FRR ( x i , θ, W FRR , φ FRR )) .
/* FRR-FLFT:
Finetuning θ with FRR according to eqn.
[Table]
D Understanding how FRR mitigates Simplicity Bias
To empirically illustrate the extent of Simplicity Bias in Neural Networks, Shah et al.
[6] introduced several synthetic and semi-synthetic datasets, where some features are explicitly simple , requiring a simpler decision boundary for prediction; while others are complex .
In this section, we demonstrate the effectiveness of the proposed Feature Reconstruction Regularizer towards mitigating Simplicity
Bias, by evaluating the same on a 10-class variant of the proposed semi-synthetic MNIST-CIFAR dataset, as discussed in the following section.
D.1 MNIST-CIFAR-10 dataset
We extend the simple binary MNIST-CIFAR dataset proposed by Shah et al.
[6] to a 10-class dataset, in order to evaluate the impact of the proposed Feature Reconstruction Regularizer in a more complex scenario when compared to the binary Colored-MNIST dataset presented in Appendix B.2.
We refer to this dataset as MNIST-CIFAR-10.
The higher complexity of this dataset allows for a more reliable evaluation of various settings such as linear probing, full network ﬁnetuning and ﬁxed-linear ﬁnetuning, with better granularity of results.
To construct this dataset, we ﬁrst deﬁne correspondences between the classes of CIFAR-10 and
MNIST.
Each image from class k of MNIST is mapped with an image from class k of CIFAR-10, with the label being set to k .
Thus, every training data sample ( x 1 , x 2 , y ) consists of x 1 and x 2 , which are images from CIFAR-10 and MNIST respectively, along with their ground truth class y .
It is to be noted that for both CIFAR-10 and MNIST, labels are always correlated with the respective images.
In such a scenario, although a classiﬁer can achieve very good performance by relying solely on the simple (MNIST) features, the goal of Out-Of-Distribution (OOD) robustness requires it to rely on the complex (CIFAR-10) features as well.
This dataset represents the toughest setting of OOD generalization, where there is no differentiation between important features and spurious correlations.
A real-world example of such a case is the classiﬁcation of swans versus bears, with the training dataset consisting of only white swans and black bears.
Here the model could either rely on shape or color for classiﬁcation.
A classiﬁcation network that relies solely on the simplest feature color, fails to generalize to the test set consisting of black swans and polar bears.
D.2 Training and Evaluation Settings
We consider two separate ResNet-18 [15] feature extractors for CIFAR-10 and MNIST respectively.
The outputs of the Global Average Pooling (GAP) layers in each of the feature extractors are concatenated to form a 1024 dimensional vector, which is given as input to the linear classiﬁer.
This architecture allows the computation of accuracy based on either a combination of both CIFARand MNIST features, or based on features of only one of the datasets.
For example, to evaluate the performance of the classiﬁer based on CIFAR-10 features alone, we replace the 512 dimensional
MNIST feature vector of each data sample with an average feature vector computed from all images in the MNIST dataset.
We refer to this as the CIFAR-AvgMNIST dataset, while the corresponding one for MNIST is refered to as the MNIST-AvgCIFAR dataset.
Similar to the work by [6], we deﬁne two additional datasets, CIFAR-RandMNIST and MNIST-RandCIFAR, where images from one of the datasets (MNIST and CIFAR-10 respectively) are randomly shufﬂed with respect to their corresponding labels.
The base training (E1, E2, E3) is done for 500 epochs, and the linear layer training / ﬁnetuning (E4 - E18) is done for 20 epochs, without any augmentations.
D.3 Experimental Results in various Training Regimes
We present the results of training on the MNIST-CIFAR-10 dataset using different algorithms in
Table 3.
The mean and standard deviation across ﬁve runs have been reported for each case.
ERM Training:
By training a randomly initialized model on the MNIST-CIFAR-10 dataset using the cross-entropy loss (E1), we obtain an accuracy of 99 .
84% on its corresponding test split.
While the accuracy of this model on the MNIST-avgCIFAR dataset is high ( 97 . 44% ), its performance on the CIFAR-avgMNIST dataset is poor ( 51 . 92% ), indicating that the model chooses to rely more on the simpler MNIST features, rather than a combination of both CIFAR and MNIST features.
[Table]
While the performance on the CIFAR-avgMNIST and MNIST-avgCIFAR datasets is sufﬁcient to understand the extent of CIFAR/ MNIST features used by the classiﬁcation head, it does not give a clear picture on the features learned by the two feature extractors.
To understand this, we reinitialize the linear classiﬁcation head randomly, and train the same using CIFAR-RandMNIST (E5) and
MNIST-RandCIFAR datasets (E6) respectively.
We obtain an accuracy of 65 . 2% on the CIFARavgMNIST dataset in the former case, indicating that although the CIFAR features learned can possibly achieve 13% higher accuracy (w.r.t. E1), the bias in the classiﬁcation head prevents them from participating in the classiﬁcation task.
The MNIST-avgCIFAR accuracy of the latter case is high as expected.
An upper bound on CIFAR-10 and MNIST accuracy that can be achieved with the selected architecture and training strategy (without using any augmentations) can be seen in E ( 88 . 53% ) and E3 ( 99 . 68% ) respectively.
Training the Linear Classiﬁcation Head: As discussed, while ERM training (E1) learns features that can be used for better OOD performance (E5), it does not effectively leverage these features for the classiﬁcation task.
We ﬁrstly explore the possibility of bridging the difference in the CIFARavgMNIST accuracy between E1 and E5 by merely retraining the linear layer.
By reinitializing and naively retraining the linear layer with Cross-entropy loss, the accuracy on CIFAR-avgMNIST improves by less than 1% (E4).
Using the proposed Feature Reconstruction Regularizer (FRR) for training the linear layer alone, the CIFAR-avgMNIST accuracy improves by 7 . 21% as shown in E8, demonstrating the effectiveness of the proposed regularizer in mitigating Simplicity Bias.
We penalize the (cid:96) ∞ norm of difference in original features and their reconstruction in addition to the minimization of cross-entropy loss.
The reconstruction based regularizer enforces the network to utilize both
CIFAR and MNIST features for classiﬁcation.
Since this regularizer resembles an orthonormality constraint on the linear classiﬁcation head, we additionally check the effectiveness of explicitly enforcing a full-rank constraint on the linear layer by minimizing the following: || WW T −
I || F (E7).
We ﬁnd that this is not effective in improving the overall accuracy, possibly because it enforces a very stringent constraint on the ﬁnal classiﬁcation layer.
Contrary to this, the proposed Feature
Reconstruction Regularizer allows more ﬂexibility by imposing this constraint only on the domain of features seen during training.
This accounts for the simple feature replication as well, enabling to view the logit layer as an information bottleneck in the reconstruction.
Finetuning (FT) and Fixed Linear Finetuning (FLFT):
We explore the ﬁnetuning of a given base model in two settings - ﬁrstly by ﬁnetuning all layers in the network (denoted as FT or FineTuning), and secondly, by freezing the parameters of the linear classiﬁcation head and ﬁnetuning only the feature extractors, which we refer to as FLFT or Fixed Linear FineTuning.
By ﬁnetuning an ERM trained base model using either of the two strategies (E9 and E10), we observe gains of less than 1% .
We observe similar gains even by ﬁnetuning the full network with FRR (E11).
Contrary to this, by using FRR-FLFT even on the ERM trained network (E12), we obtain 7 .
29% improvement over the base model.
This shows that, by allowing the full network to change while imposing the FRR constraint, the network can continue to rely on simple features, possibly by reducing the number of complex features learned by the feature extractor.
However, by freezing the weights of the linear layer and further imposing this constraint, the network is forced to reﬁne the CIFAR features that are already being used for prediction.
Combining FRR-L and FRR-FLFT: While we obtain similar order of gains ( ∼ 7% ) using both
FRR-L and FRR-FLFT
individually, the former improves the diversity of features being considered by the classiﬁcation head, while the latter improves the quality of the features themselves.
We therefore propose a training strategy that combines both FRR-L and FRR-FLFT.
Using this, we obtain gains of 16 .
2% over the ERM baseline as shown in E16, indicating that the combination of FRR-L and
FRR-FLFT has a compounding effect by ﬁrstly selecting diverse features, and further reﬁning these features to be more useful for classiﬁcation.
Although FRR-L followed by FRR-FT (E15) is also effective, it has about 6% lesser gains when compared to the proposed approach of FRR-L + FRRFLFT.
We note that following up FRR-L with ERM-FT (E13) or ERM-FLFT
(E14) also reﬁnes the learned features, making them more suitable for the classiﬁcation task, yielding 2 .
6% and 4 .
6% gains respectively over FRR-L.
We verify the quality of features learned by the feature extractors after the proposed training strategy
FRR-L + FRR-FLFT by reinitializing and retraining the linear classiﬁer on CIFAR-RandMNIST (E17) and MNIST-RandCIFAR (E18) datasets respectively.
We observe considerable gains of around 15% on MNIST-CIFAR-10 accuracy using CIFAR-RandMNIST training when compared to ERM (E5), demonstrating that the proposed approach not only results in more CIFAR features being used for classiﬁcation, but also leads to the learning of better CIFAR features.
E Details on the OOD Generalization setting considered
The problem of improving robustness to distribution shifts has been studied in several settings, where, in addition to labeled source domain data, varying levels of access to the target domain data is assumed.
Some of the well-researched settings include - Unsupervised Domain Adaptation , with access to only unlabeled target domain data [16, 17], and Domain Generalization , where typically data from several source distributions is assumed to be available, and the target domain in unseen during training [18, 19, 11]).
In the latter case, it is assumed that all training data samples are annotated with domain labels as well, so that training algorithms can explicitly impose invariance to attributes that cause a distribution shift in input data without change in their label distribution [20, 17, 21, 22, 23].
A more challenging case is when the training data belongs to several distributions that may not even be sufﬁciently discernable to have explicit domain annotations, or may contain multidimensional distribution shifts, such as weather, time of the day and geographical location, that cannot be easily annotated or clustered.
We investigate this crucial setting which has been relatively less researched, and refer to it as Aggregated Domain Generalization , as introduced by [24].
We note that this setting is different from the case of training on data from a single domain such as ImageNet, and evaluating on distribution shifts [25], due to the availability of an aggregate of source domains during training, which can enable the effective use of in-domain validation set for hyperparameter selection.
While there have been several approaches to improve the performance of models in the setting of Domain Generalization, Gulrajani et al.
[11] show that when evaluated fairly, that is, without assuming access to the test domain data even for selecting the best set of hyperparameters, none of the approaches perform consistently better than standard training using Empirical Risk Minimization (ERM) .
Furthermore, we consider the setting of Aggregated Domain Generalization , which is more challenging due to the absence of domain labels during both training and validation.
F Related Works
Learning diverse classiﬁers to counter simplicity bias: Recent works have shown that ERM trained models learn diverse features, however, the linear layer fails at capturing and utilizing these diverse features properly.
There have been several attempts at training classiﬁers which can make use of such diverse features.
[26] train a number of linear classiﬁers on top of a pre-trained network with a diversity regularizer, which encourages the classiﬁers to rely on different features.
[27] extend this idea further to train a classiﬁer which is “orthogonal" to the original network, i.e. its predictions do not depend on the same features that the original network’s predictions depend on.
However, this assumes access to an unbiased network, whose predictions already take both sets of features into account.
[28] show that reweighting train set examples and retraining the last layer of a pre-trained deep network can alleviate spurious correlations, provided one can access a balanced dataset.
In contrast to these methods, our method can work simply on the training set data, and produce a single classiﬁer which is debiased.
[29] propose to mute the features with highest gradients, and use only the other features to make a prediction.
While this method suppresses the maximally used features, it does not encourage the learning of hard-to-learn features, which is directly realized using our loss formulation.
[10] suggest that ﬁnetuning the ﬁnal linear layer ﬁrst before ﬁnetuning the entire network can make it more robust to OOD shifts, and we utilize this insight in the FRR-FLFT phase of our method.
A complementary approach to this problem is to learn features that are more diverse [30].
We note that applying our proposed method on top of such techniques would encourage the classiﬁer to use the diverse features effectively, and this can further beneﬁt the performance.
Domain Generalization and OOD robustness: The performance of neural networks is known to drop when there is a mismatch in the train and test distributions [25], and methods to mitigate this have been gaining a lot of attention in recent years.
The problem has been studied under various assumptions on distribution shift.
The commonly studied setting of domain generalization [11, 19] assumes that the train distribution consists of a mixture of distinct distributions (called domains), with each train sample having a domain label associated with it.
The stronger setting of aggregate domain generalization [24, 31] assumes training data to be drawn from a mixture of distributions, but does not assume the availability of domain labels.
Finally, OOD robustness [25, 32] drops all of these assumptions.
Most works tackling the domain generalization problem attempt to train a model whose predictions are invariant to the domain label [19, 22], or try to align the features of the model for examples from different domains [23, 33].
However, since we aim to tackle the stronger setting of OOD generalization, we do not use domain labels.
Tackling the OOD robustness problem, [24] and [31] ﬁrst cluster training examples into “pseudo-domains", after which standard domain generalization techniques are used.
Another recent line of works propose using model averaging [12, 34] and/or ensembling [14] for better OOD generalization.
These techniques are complementary to our contribution, and we demonstrate how they can beneﬁt each other in our empirical evaluation.
G Experimental Details on DomainBed
We test our approach on the DomainBed benchmark [11] comprising of ﬁve different datasets, each of which have k domains.
For each dataset, we train a model on k − 1 domains, and test it on the left out domain.
The average out-of-domain performance across the k held-out domains is then reported.
In this section we describe the hyper-parameter selection strategy and the ranges considered for our approach.
In line with the DomainBed testbench, we use ImageNet pretrained ResNet-50 models for all algorithms.
We use random search to select hyperparameters for our algorithm, and use the suggested hyperparameters for the other baselines.
We train for 3000 (5000 for DomainNet) steps in the FRR-L phase, and 5000 (10000 for DomainNet) steps in the FRR-FLFT phase.
The batch size is ﬁxed to 32, and SWAD hyper-parameters are the same as those used by [12].
We use the indomain accuracy protocol from [11] to select hyper-parameters for each domain of each dataset, and search over 8 random combinations of hyper-parameters for each.
The range of the hyperparameters is shown in Table 4.
Note that we experiment with two implementations of (cid:96) ∞ norm: (cid:96) 1 , ∞ , where we ﬁrst compute the (cid:96) ∞ of feature reconstruction for each example in a batch and then average it across the batch, and (cid:96) ∞ , 1 where we compute the average (cid:96) 1 reconstruction norm of each feature across the batch, and then apply (cid:96) ∞ norm on this m dimensional vector.
All our experiments were done on single V100 GPUs.
[Table]
H Ablations on DomainBed
Comparing the choices for φ : In Table 5, we experiment with various architectures for the decoder φ when computing FRR according to equation 1.
We consider using a two layer neural network as the decoder φ (FRR-LDeeper), and also consider setting φ = W T (FRR-LShared), i [Table] e. explicitly tying the weights of the decoder and the classiﬁer layer.
Overall, both these variants are worse than the default single layer, free parameterization of φ .
We believe that this happens because the latter approach enforces a much stricter constraint on W , leading to poorer in-domain accuracy, while the former approach enforces a weaker constraint, potentially enabling reconstruction of more complex features from a smaller amount of information about them in the logits.
Both these have a detrimental effect on the overall performance of the model
[Table]
[Table]
I Domain wise accuracies
In this section
[Table]
we show detailed results of Table 1 in the main text
[Table]
[Table] [Table]
[Table] [Table] [Table]