Original: F EATURE R ECONSTRUCTION F ROM O UTPUTS C AN M ITIGATE S IMPLICITY B IAS IN N EURAL 
Revised: F EATURE R ECONSTRUCTION F ROM O UTPUTS C AN M ITIGATE S IMPLICITY B IAS IN N EURAL N ETWORKS

Original: N

Original: ETWORKS

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Deep Neural Networks are known to be brittle to even minor distribution shifts compared to the training distribution.
Revised: Deep Neural Networks are known to be brittle to even minor distribution shifts compared to the training distribution.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Namely, simple features are replicated several times in the learned representations while complex features might not be replicated.
Revised: Namely, simple features are replicated several times in the learned representations while complex features might not be replicated.

Original: 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.
Revised: 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.

Original: To mitigate this bias, we propose Feature Reconstruction Regularizer (FRR) to ensure that the learned features can be reconstructed back from the logits.
Revised: To mitigate this bias, we propose Feature Reconstruction Regularizer (FRR) to ensure that the learned features can be reconstructed back from the logits.

Original: The use of FRR in linear layer training ( FRR-L ) encourages the use of more diverse features for classiﬁcation.
Revised: The use of FRR in linear layer training ( FRR-L ) encourages the use of more diverse features for classiﬁcation.

Original: 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.
Revised: 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.

Original: Using this simple solution, we demonstrate up to 15% gains in OOD accuracy on the recently introduced semi-synthetic datasets with extreme distribution shifts.
Revised: Using this simple solution, we demonstrate up to 15% gains in OOD accuracy on the recently introduced semi-synthetic datasets with extreme distribution shifts.

Original: Moreover, we demonstrate noteworthy gains over existing SOTA methods on the standard OOD benchmark DomainBed as well.
Revised: Moreover, we demonstrate noteworthy gains over existing SOTA methods on the standard OOD benchmark DomainBed as well

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Despite the remarkable success of Deep Neural Networks (DNNs) in various ﬁelds, they are known to be brittle against even minor shifts in the data distribution during inference, which are not uncommon in a real world setting (Quinonero-Candela et al., 2008; Torralba & Efros, 2011).
Revised: Despite the remarkable success of Deep Neural Networks (DNNs) in various ﬁelds, they are known to be brittle against even minor shifts in the data distribution during inference, which are not uncommon in a real world setting (Quinonero-Candela et al., 2008; Torralba & Efros, 2011).

Original: For example, a self-driving car that works well in normal weather may perform poorly when it is snowing, leading to disastrous outcomes.
Revised: For example, a self-driving car that works well in normal weather may perform poorly when it is snowing, leading to disastrous outcomes.

Original: The need for improving the robustness of such systems against distribution shifts has sparked interest in the area of Out-Of-Distribution or OOD generalization (Hendrycks &
Revised: The need for improving the robustness of such systems against distribution shifts has sparked interest in the area of Out-Of-Distribution or OOD generalization (Hendrycks &

Original: Dietterich, 2019; Gulrajani & Lopez-Paz, 2020).
Revised: Dietterich, 2019; Gulrajani & Lopez-Paz, 2020).

Original: In this work, we aim to tackle the problem of OOD generalization of Neural Networks in a covariateshift (Shimodaira, 2000) 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.
Revised: In this work, we aim to tackle the problem of OOD generalization of Neural Networks in a covariateshift (Shimodaira, 2000) 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.

Original: 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 (Arpit et al., 2017; Valle-Perez et al., 2018).
Revised: 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 (Arpit et al., 2017; Valle-Perez et al., 2018).

Original: 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 (Shah et al., 2020).
Revised: 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 (Shah et al., 2020).

Original: The extent of Simplicity Bias seen in models is a result of two important factors - diversity of features learned by the feature extractor 1 , and the extent to which these diverse features are used for the task at hand, such as classiﬁcation 2 .
Revised: The extent of Simplicity Bias seen in models is a result of two important factors - diversity of features learned by the feature extractor 1 , and the extent to which these diverse features are used for the task at hand, such as classiﬁcation 2 .

Original: 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 (Rosenfeld et al., 2022; Kirichenko et al., 2022b).
Revised: 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 (Rosenfeld et al., 2022; Kirichenko et al., 2022b).

Original: Does this imply that brittleness of models can be attributed to the learning of the classiﬁcation head alone?
Revised: Does this imply that brittleness of models can be attributed to the learning of the classiﬁcation head alone?

Original: 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 (Soudry et al., 2018)?
Revised: 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 (Soudry et al., 2018)?

Original: 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.
Revised: 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.

Original: 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 .
Revised: 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 .

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: As shown in Fig.2, we ﬁrstly propose to train the linear classiﬁer alone by freezing the weights of the feature extractor.
Revised: As shown in Fig.2, we ﬁrstly propose to train the linear classiﬁer alone by freezing the weights of the feature extractor.

Original: This formulation enables the learning of an Invertible Mapping in the output layer, speciﬁcally for the domain of features seen during training.
Revised: This formulation enables the learning of an Invertible Mapping in the output layer, speciﬁcally for the domain of features seen during training.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Consequently, the learnt linear classiﬁer also gives more importance to non replicated complex features while making predictions.
Revised: Consequently, the learnt linear classiﬁer also gives more importance to non replicated complex features while making predictions.

Original: 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.
Revised: 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.

Original: In order to do this, we freeze the linear classiﬁcation head, and further ﬁnetune the backbone with FRR.
Revised: In order to do this, we freeze the linear classiﬁcation head, and further ﬁnetune the backbone with FRR.

Original: We ﬁnd that this encourages the network to indeed learn better quality features that are more relevant for classiﬁcation.
Revised: We ﬁnd that this encourages the network to indeed learn better quality features that are more relevant for classiﬁcation.

Original: We list the key contributions of this work below -
Revised: We list the key contributions of this work below -

Original: • Key Observation : We provide a crisp hypothesis of “feature replication” to explain the brittleness of ERM trained neural networks to OOD data (Sec 3.1).
Revised: • Key Observation : We provide a crisp hypothesis of “feature replication” to explain the brittleness of ERM trained neural networks to OOD data (Sec 3.1).

Original: Using this, we further provide theoretical and empirical evidence to justify the existence of Simplicity Bias in maximum margin classiﬁers.
Revised: Using this, we further provide theoretical and empirical evidence to justify the existence of Simplicity Bias in maximum margin classiﬁers.

Original: • Novel Algorithm based on the Observation : Based on this, we introduce a novel FRR regularizer to safeguard against the feature replication phenomenon (Sec 3.2).
Revised: • Novel Algorithm based on the Observation : Based on this, we introduce a novel FRR regularizer to safeguard against the feature replication phenomenon (Sec 3.2).

Original: We also provide theoretical support for FRR in an intuitive data distribution setting.
Revised: We also provide theoretical support for FRR in an intuitive data distribution setting.

Original: 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 4).
Revised: 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 4).

Original: • 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 2) 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 for OOD generalization(Table 3).
Revised: • 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 2) 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 for OOD generalization(Table 3).

Original:  R ELATED W ORKS
Revised:  R ELATED W ORKS

Original: 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.
Revised: 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.

Original: There have been several attempts at training classiﬁers which can make use of such diverse features.
Revised: There have been several attempts at training classiﬁers which can make use of such diverse features.

Original: Teney et al.
Revised: Teney et al.

Original: (2022) train a number of linear classiﬁers on top of a pretrained network with a diversity regularizer, which encourages the classiﬁers to rely on different features.
Revised: (2022) 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.

Original: Xu et al.
Revised: Xu et al.

Original: and Bahng et al.
Revised: and Bahng et al.

Original: (2020) propose to train debiased classiﬁers which are statistically independent from trained biased networks, but these need careful design and prior knowledge of the biases in trained networks.
Revised: (2020) propose to train debiased classiﬁers which are statistically independent from trained biased networks, but these need careful design and prior knowledge of the biases in trained networks.

Original: Kirichenko et al.
Revised: Kirichenko et al.

Original: (2022a) 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.
Revised: (2022a) 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.

Original: In contrast to these methods, our method can work simplythe on the training set data, and produce a single classiﬁer which is debiased.
Revised: In contrast to these methods, our method can work simply on the training set data, and produce a single classiﬁer which is debiased.

Original: Huang et al.
Revised: Huang et al.

Original: (2020) propose to mute the features with highest gradients, and use only the other features to make a prediction.
Revised: (2020) propose to mute the features with highest gradients, and use only the other features to make a prediction.

Original: 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.
Revised: While this method suppresses the maximally used features, it does not encourage the learning of hard-tolearn features, which is directly realized using our loss formulation.

Original: Kumar et al.
Revised: Kumar et al.

Original: (2022) 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.
Revised: (2022) 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.

Original: A complementary approach to this problem is to learn features that are more diverse (Zhang et al., 2022).
Revised: A complementary approach to this problem is to learn features that are more diverse (Zhang et al., 2022;

Original: For example, Wang et al.
Revised: Wang et al., 2019).

Original: (2019) propose to regularize local features learnt by the network to avoid being predictive of the label, so that the network is forced to rely on the global information in the image.

Original: 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.
Revised: 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.

Original: 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 (Hendrycks & Dietterich, 2019), and methods to mitigate this have been gaining a lot of attention in recent years.
Revised: 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 (Hendrycks & Dietterich, 2019), and methods to mitigate this have been gaining a lot of attention in recent years.

Original: The problem has been studied under various assumptions on distribution shift.
Revised: The problem has been studied under various assumptions on distribution shift.

Original: The commonly studied setting of domain generalization (Gulrajani & Lopez-Paz, 2020; Li et al., 2018a) 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.
Revised: The commonly studied setting of domain generalization (Gulrajani & Lopez-Paz, 2020; Li et al., 2018a) 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.

Original: The stronger setting of aggregate domain generalization (Thomas et al., 2021; Matsuura & Harada, 2020) assumes training data to be drawn from a mixture of distributions, but does not assume the availability of domain labels.
Revised: The stronger setting of aggregate domain generalization (Thomas et al., 2021; Matsuura & Harada, 2020) assumes training data to be drawn from a mixture of distributions, but does not assume the availability of domain labels.

Original: Finally, OOD robustness (Hendrycks & Dietterich, 2019; Koh et al., 2021) drops all of these assumptions.
Revised: Finally, OOD robustness (Hendrycks & Dietterich, 2019; Koh et al., 2021) drops all of these assumptions.

Original: Most works tackling the domain generalization problem attempt to train a model whose predictions are invariant to the domain label (Li et al., 2018a; Arjovsky et al., 2019), or try to align the features of the model for examples from different domains (Shi et al., 2021; Shankar et al., 2018).
Revised: Most works tackling the domain generalization problem attempt to train a model whose predictions are invariant to the domain label (Li et al., 2018a; Arjovsky et al., 2019), or try to align the features of the model for examples from different domains (Shi et al., 2021; Shankar et al., 2018).

Original: However, since we aim to tackle the stronger setting of OOD generalization, we do not use domain labels.
Revised: However, since we aim to tackle the stronger setting of OOD generalization, we do not use domain labels.

Original: Tackling the OOD robustness problem, Thomas et al.
Revised: Tackling the OOD robustness problem, Thomas et al.

Original: and Matsuura & Harada (2020)
Revised: and Matsuura & Harada (2020)

Original: ﬁrst cluster training examples into “pseudo-domains”, after which standard domain generalization techniques are used.
Revised: ﬁrst cluster training examples into “pseudo-domains”, after which standard domain generalization techniques are used.

Original: Another recent line of works propose using model averaging (Cha et al., 2021;
Revised: Another recent line of works propose using model averaging (Cha et al., 2021;

Original: Li et al., 2022) and/or ensembling (Arpit et al., 2021) for better OOD generalization.
Revised: Li et al., 2022) and/or ensembling (Arpit et al., 2021) for better OOD generalization.

Original: These techniques are complementary to our contribution, and we demonstrate how they can beneﬁt each other in our empirical evaluation.
Revised: These techniques are complementary to our contribution, and we demonstrate how they can beneﬁt each other in our empirical evaluation.

Original:  F EATURE R EPLICATION H YPOTHESIS
Revised:  F EATURE R EPLICATION H YPOTHESIS

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Hence, the outputs of the network are brittle to distribution shifts that change such replicated features.
Revised: Hence, the outputs of the network are brittle to distribution shifts that change such replicated features.

Original: In this section, we provide empirical and theoretical evidence for FRH, and propose a new regularizer – FRR– to mitigate this effect.
Revised: In this section, we provide empirical and theoretical evidence for FRH, and propose a new regularizer – FRR– to mitigate this effect.

Original: We ﬁrst introduce some useful notations.
Revised: We ﬁrst introduce some useful notations.

Original: Let f θ ( x ) : R d → R m be the feature extractor of a neural network parameterized by weights θ , and W ∈ R m × k be the weight matrix of the linear classiﬁer.
Revised: Let f θ ( x ) : R d → R m be the feature extractor of a neural network parameterized by weights θ , and W ∈ R m × k be the weight matrix of the linear classiﬁer.

Original: For input x ∈ R d , the output of the network is W T f θ ( x ) ∈ R k .
Revised: For input x ∈ R d , the output of the network is W T f θ ( x ) ∈ R k .

Original: 3. E MPIRICAL VALIDATION OF F EATURE R EPLICATION H YPOTHESIS 
Revised: 3. E MPIRICAL VALIDATION OF F EATURE R EPLICATION H YPOTHESIS (FRH) IN ERM

Original: (FRH) IN ERM

Original: [Table]

Original: Coloured MNIST dataset.
Revised: Coloured MNIST dataset.

Original: To empirically demonstrate feature replication, we use a binarized version of the coloured MNIST dataset (Gulrajani & Lopez-Paz, 2020).
Revised: To empirically demonstrate feature replication, we use a binarized version of the coloured MNIST dataset (Gulrajani & Lopez-Paz, 2020).

Original: We construct this dataset by ﬁrst assigning two digits of the MNIST dataset, namely “1” and “5”, to classes 0 and 1 respectively.
Revised: We construct this dataset by ﬁrst assigning two digits of the MNIST dataset, namely “1” and “5”, to classes 0 and 1 respectively.

Original: While training the network, we super-impose images of “1” onto colours of range R 

Revised: While training the network, we super-impose images of “1” onto colours of range R

Original: = [(115 , 0 , 0) − (256 , 141 , 0)] (i.e. red), and images of “5” onto colours of range R 1 =
Revised: = [(115 , 0 , 0) − (256 , 141 , 0)] (i.e. red), and images of “5” onto colours of range R 1 =

Original: [(0 , 115 , 0) − (141 , 256 , 0)]
Revised: [(0 , 115 , 0) − (141 , 256 , 0)]

Original: (i.e. green).
Revised: (i.e. green).

Original: The dataset is constructed such that the simple feature, namely colour, is weakly correlated with the labels, while the complex shape features are strongly correlated with labels.
Revised: The dataset is constructed such that the simple feature, namely colour, is weakly correlated with the labels, while the complex shape features are strongly correlated with labels.

Original: See Appendix D.1 for more details about the dataset.
Revised: See Appendix D.1 for more details about the dataset.

Original:  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.
Revised: [Table] 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.

Original: 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.
Revised: 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.

Original: We train a four layered CNN on this data.
Revised: We train a four layered CNN on this data.

Original: 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.
Revised: 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.

Original: We also compute the correlation of these features with the output of the network ( W T f θ ( x ) ) over inputs from the test domain.
Revised: We also compute the correlation of these features with the output of the network ( W T f θ ( x ) ) over inputs from the test domain.

Original: This gives us information of the learnt features, and their contributions to the ﬁnal prediction of the network.
Revised: This gives us information of the learnt features, and their contributions to the ﬁnal prediction of the network.

Original: Note that the feature dimension is m = 32 , and the output dimension is k = 1 .
Revised: Note that the feature dimension is m = 32 , and the output dimension is k = 1 .

Original: Observations :
Revised: Observations :

Original: In Table 1, we report the number of colour features, shape features, and the average correlation of each of these with the ﬁnal prediction.
Revised: In Table 1, we report the number of colour features, shape features, and the average correlation of each of these with the ﬁnal prediction.

Original: 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 ourFeature
Revised: 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

Original:  Replication Hypothesis .
Revised: Feature Replication Hypothesis .

Original: We also visualize the inter-feature correlation of the learnt features in Fig 5, which shows blocks of highly correlated features, further validating our hypothesis.
Revised: We also visualize the inter-feature correlation of the learnt features in Fig 5, which shows blocks of highly correlated features, further validating our hypothesis.

Original: Furthermore, we note that correlation of the output with the shape features is lower, leading to OOD accuracy of 59% only.
Revised: We note that correlation of the output with the shape features is lower, leading to OOD accuracy of 59% .

Original: 3. F EATURE R ECONSTRUCTION R EGULARIZER (FRR)
Revised: 3. F EATURE R ECONSTRUCTION R EGULARIZER (FRR)

Original: To alleviate simple feature replication issue, we propose Feature Reconstruction Regularizer (FRR) to enforce that the learned features can be reconstructed from the output logits.
Revised: To alleviate simple feature replication issue, we propose Feature Reconstruction Regularizer (FRR) to enforce that the learned features can be reconstructed from the output logits.

Original: We propose to retrain the ﬁnal linear layer using this regularizer to allow the model to utilize diverse features to compute the ﬁnal output.
Revised: We propose to retrain the ﬁnal linear layer using this regularizer to allow the model to utilize diverse features to compute the ﬁnal output.

Original: We implement this by introducing another neural network with the objective of reconstructing the features of the network from the output logits, i.e. features f θ ( x ) should be recoverable from the predictions of the network through a transform T φ ( . ) parameterized by φ .
Revised: We implement this by introducing another neural network with the objective of reconstructing the features of the network from the output logits, i.e. features f θ ( x ) should be recoverable from the predictions of the network through a transform T φ ( . ) parameterized by φ .

Original: That is, FRR is given by:
Revised: That is, FRR is given by:

Original: [Equation]
Revised: [Equation]

Original: where || .
Revised: where || .

Original: ||
Revised: ||

Original: p denotes the (cid:96) p norm.
Revised: p denotes the (cid:96) p norm.

Original: We set this norm to be (cid:96) ∞ or (cid:96) 1 in our experiments.
Revised: We set this norm to be (cid:96) ∞ or (cid:96) 1 in our experiments.

Original: In the simplest case, T φ ( y ) =
Revised: In the simplest case, T φ ( y ) =

Original: φy , where φ ∈ R m × k .
Revised: φy , where φ ∈ R m × k .

Original: Note that in order to ﬁnd the appropriate φ , we jointly optimize W and φ using gradient descent based optimizers.
Revised: Note that in order to ﬁnd the appropriate φ , we jointly optimize W and φ using gradient descent based optimizers.

Original: We also experiment with φ being a more complex neural network.
Revised: We also experiment with φ being a more complex neural network.

Original: We empirically validate FRR on Coloured MNIST , where using FRR with the linear layer leads to lower correlation with Colour compared to standard ERM (Table 1).
Revised: We empirically validate FRR on Coloured MNIST , where using FRR with the linear layer leads to lower correlation with Colour compared to standard ERM (Table 1).

Original: Consequently, OOD accuracy improves by 5% over ERM.
Revised: Consequently, OOD accuracy improves by 5% over ERM.

Original: 3. FRH & FRR: T HEORETICAL A NALYSIS
Revised: 3. FRH & FRR: T HEORETICAL A NALYSIS

Original: We now present a simple and intuitive data distribution with feature replication that highlights the
Revised: We now present a simple and intuitive data distribution with feature replication that highlights the

Original: OOD brittleness of standard ERM, and also demonstrates FRR can be signiﬁcantly more robust.
Revised: OOD brittleness of standard ERM, and also demonstrates FRR can be signiﬁcantly more robust.

Original: Data Distribution : Consider a linearly separable distribution consisting of two factors of variation as shown in Figure 1.
Revised: Data Distribution : Consider a linearly separable distribution consisting of two factors of variation as shown in Figure 1.

Original: That is, consider the following distribution ( x, y ) ∼ D , where, y = ± 1 with probability 0 .
Revised: That is, consider the following distribution ( x, y ) ∼ D , where, y = ± 1 with probability 0 .

Original: [Equation]
Revised: [Equation]

Original: We consider d (= 0 or 5) feature replications (Rep) along the  to a  solution in the non-replicated  space, resulting in improved OOD robustness (
Revised: We consider d (= 0 or 5) feature replications (Rep) along the y- FRR converges to a maximum margin solution in the non-replicated feature space, resulting in improved OOD robustness (d)

Original: 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,
Revised: 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,

Original: [Equation]
Revised: [Equation]

Original: i.e., x 1 is repeated d times.
Revised: i.e., x 1 is repeated d times.

Original: The joint distribution of features and labels is denoted by ˜ D .
Revised: The joint distribution of features and labels is denoted by ˜ D .

Original: Finally, we deﬁne the l 2 max margin classiﬁer over a distribution D as w MM :=
Revised: Finally, we deﬁne the l 2 max margin classiﬁer over a distribution D as w MM :=

Original: arg min w 12 (cid:107) w (cid:107) 22 subject to y · (cid:104) w, x (cid:105) ≥ 1 ∀ ( x, y
Revised: arg min w 12 (cid:107) w (cid:107) 22 subject to y · (cid:104) w, x (cid:105) ≥ 1 ∀ ( x, y

Original: ) ∈
Revised: ) ∈

Original: Supp ( D ) .
Revised: Supp ( D ) .

Original: Then we have the following results:
Revised: Then we have the following results:

Original: Claim 3.1 (Brittleness due to Feature Replication) .
Revised: Claim 3.1 (Brittleness due to Feature Replication) .

Original: Consider the data distribution given in Equation 2, 3.
Revised: Consider the data distribution given in Equation 2, 3.

Original: Then, the following holds: (1.)
Revised: Then, the following holds: (1.)

Original: The max-margin classiﬁer w MM over D is given by w MM =
Revised: The max-margin classiﬁer w MM over D is given by w MM =

Original: [1 , 1] , and (2.)
Revised: [1 , 1] , and (2.)

Original: 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 .
Revised: 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 .

Original: 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.
Revised: 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.

Original: Hence, even a minor change in this replicated feature in the input space would be ampliﬁed in the output of the classiﬁer.
Revised: Hence, even a minor change in this replicated feature in the input space would be ampliﬁed in the output of the classiﬁer.

Original: This is especially concerning in light of the observations in Table 1, which validate the Feature Replication Hypothesis in Coloured MNIST.
Revised: This is especially concerning in light of the observations in Table 1, which validate the Feature Replication Hypothesis in Coloured MNIST.

Original: Proposition 3.2 (Robustness of FRR) .
Revised: Proposition 3.2 (Robustness of FRR) .

Original: Denote the average feature reconstruction loss
Revised: Denote the average feature reconstruction loss

Original: L FRR ( ˜ w, ˜ φ ) :=max
Revised: L FRR ( ˜ w, ˜ φ ) :=

Original: [Equation]
Revised: [Equation]

Original: ≤ i ≤ d +
Revised: E

Original: [Equation]
Revised: [Equation]

Original: E (
Revised: x,y

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Original: x,y ) ∼
Revised: ˜ w, ˜ x

Original: [Equation]
Revised: [Equation]

Original: D (cid:104) ( (cid:104) ˜

Original: w, ˜ x (cid:105) ˜

Original:  φ i − ˜ x i ) 2 (cid:105) and consider any ( ˜ w ∗ , ˜ φ ∗ ) satisfying:[Equation] [Equation] φ
Revised: ˜ φ i − ˜ x i ) 2 (cid:105) and consider any ( ˜ w ∗ , ˜ φ ∗ ) satisfying:

Original: [Equation]
Revised: [Equation]

Original: ∈ arg min ( ˜ [Equation] , ˜ φ ) L FRR (
Revised: subject to

Original: [Equation]
Revised: [Equation]

Original: w, ˜ φ ) subject to y · (cid:104) ˜ w, ˜ x (cid:105) ≥ 0 ∀ (˜ x, y

Original: ) ∈ Supp (cid:16) ˜ D (cid:17) .

Original: We have that: ˜ w ∗ 1 + · · · + ˜ w ∗ d = ˜ w ∗ d +1 .
Revised: We have that: ˜ w ∗ 1 + · · · + ˜ w ∗ d = ˜ w ∗ d +1 .

Original: Consequently, we have (cid:104) ˜
Revised: Consequently, we have (cid:104) ˜

Original: w ∗ , ˜ x (cid:105) ∝ (cid:104)
Revised: w ∗ , ˜ x (cid:105) ∝ (cid:104)

Original: w MM , x (cid:105) for all x ∈ R 2 .
Revised: w MM , x (cid:105) for all x ∈ R 2 .

Original: Practically, we can implement the above as (cid:96) 2 , ∞ over a batch.
Revised: Practically, we can implement the above as (cid:96) 2 , ∞ over a batch.

Original: 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.
Revised: 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.

Original: This is equivalent to a maximum margin classiﬁer in the non-replicated feature space , thereby resulting in enhanced robustness to distribution shifts.
Revised: This is equivalent to a maximum margin classiﬁer in the non-replicated feature space , thereby resulting in enhanced robustness to distribution shifts.

Original: Same is reﬂected in
Revised: Same is reﬂected in

Original: Figure 1 (c), (d) which show impact of FRR on the trained boundary in the non-replicated feature space.
Revised: Figure 1 (c), (d) which show impact of FRR on the trained boundary in the non-replicated feature space.

Original: We defer the proofs of the above to Appendix A. 
Revised: We defer the proofs of the above to Appendix A. We also provide a more general result by assuming correlated feature representations in Appendix A.

Original:  T RAINING P ROCEDURE
Revised:  T RAINING P ROCEDURE

Original: Pretraining :
Revised: Pretraining :

Original: In order to learn features which are relevant to the train distribution, we ﬁrst pretrain our model using standard ERM with the cross-entropy loss L std ( W, θ, ( x, y )) .
Revised: In order to learn features which are relevant to the train distribution, we ﬁrst pretrain our model using standard ERM with the cross-entropy loss L std ( W, θ, ( x, y )) .

Original: FRR-L : Since ERM training is known to learn several rich and diverse features, we freeze the backbone parameters θ , and retrain the ﬁnal layer W as the following-
Revised: FRR-L : Since ERM training is known to learn several rich and diverse features, we freeze the backbone parameters θ , and retrain the ﬁnal layer W as the following-

Original: [Equation]
Revised: [Equation]

Original: where λ L is a hyperparameter weighing the two losses.
Revised: where λ L is a hyperparameter weighing the two losses.

Original: We train W and φ jointly.
Revised: We train W and φ jointly.

Original: We refer to this step as FRR-L, i.e. Feature Reconstruction Regularizer - Linear, since we only train the linear layer.
Revised: We refer to this step as FRR-L, i.e. Feature Reconstruction Regularizer - Linear, since we only train the linear layer.

Original: [Figure]
Revised: [Figure]

Original: FRR-FLFT :
Revised: FRR-FLFT :

Original: Following the suggestions of Kumar et al.
Revised: Following the suggestions of Kumar et al.

Original: (2022), we follow up the linear layer training with the ﬁnetuning of the feature extractor θ with a weighted combination of the crossentropy loss and FRR, weighted by a hyper-parameter λ FLFT .
Revised: (2022), we follow up the linear layer training with the ﬁnetuning of the feature extractor θ with a weighted combination of the crossentropy loss and FRR, weighted by a hyper-parameter λ FLFT .

Original: In this step we freeze the weights of the linear layer to improve the stability of training.
Revised: In this step we freeze the weights of the linear layer to improve the stability of training.

Original: We do this since naively using this constraint during network training could amplify the Simplicity Bias in networks in degenerate cases.
Revised: We do this since naively using this constraint during network training could amplify the Simplicity Bias in networks in degenerate cases.

Original: For example, the backbone could learn to output a single replicated simple feature which is predictive enough on the training data.
Revised: For example, the backbone could learn to output a single replicated simple feature which is predictive enough on the training data.

Original: Reconstructing such a feature from logits would also be easy, but such a network might not generalize well.
Revised: Reconstructing such a feature from logits would also be easy, but such a network might not generalize well.

Original: Formally, the optimization problem for this step is -
Revised: Formally, the optimization problem for this step is -

Original: [Equation]
Revised: [Equation]

Original: We view this step as “sharpening” the features for more accurate predictions.
Revised: We view this step as “sharpening” the features for more accurate predictions.

Original: Freezing the linear head makes sure that the features do not collapse to a degenerate solution.
Revised: Freezing the linear head makes sure that the features do not collapse to a degenerate solution.

Original: Our training algorithm is summarized in Algorithm-1 and the training pipeline is illustrated in Figure 2.
Revised: Our training algorithm is summarized in Algorithm-1 and the training pipeline is illustrated in Figure 2.

Original:  E XPERIMENTAL R ESULTS 5. U NDERSTANDING HOW FRR MITIGATES S IMPLICITY B IAS
Revised:  E XPERIMENTAL R ESULTS 5. U NDERSTANDING HOW FRR MITIGATES S IMPLICITY B IAS

Original: To empirically illustrate the extent of Simplicity Bias in Neural Networks, Shah et al.
Revised: To empirically illustrate the extent of Simplicity Bias in Neural Networks, Shah et al.

Original: (2020) introduced several synthetic and semi-synthetic datasets, where some features are explicitly simple , requiring a simpler decision boundary for prediction; while others are complex .
Revised: (2020) introduced several synthetic and semi-synthetic datasets, where some features are explicitly simple , requiring a simpler decision boundary for prediction; while others are complex .

Original: 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
Revised: 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

Original: MNIST-CIFAR dataset, as discussed in the following section.
Revised: MNIST-CIFAR dataset, as discussed in the following section.

Original: 5.1. MNIST-CIFAR-10 DATASET
Revised: 5.1. MNIST-CIFAR-10 DATASET

Original: We extend the simple binary MNIST-CIFAR dataset proposed by Shah et al.
Revised: We extend the simple binary MNIST-CIFAR dataset proposed by Shah et al.

Original: 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 Section3.
Revised: 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 Section3.

Original: We refer to this dataset as MNIST-CIFAR-10.
Revised: We refer to this dataset as MNIST-CIFAR-10.

Original: The higher complexity of this dataset allows for a more reliable evaluation of various settings such as linear probing, full network ﬁnetuning and ﬁxedlinear ﬁnetuning, with better granularity of results.
Revised: The higher complexity of this dataset allows for a more reliable evaluation of various settings such as linear probing, full network ﬁnetuning and ﬁxedlinear ﬁnetuning, with better granularity of results.

Original: To construct this dataset, we ﬁrst deﬁne correspondences between the classes of CIFAR-10 and
Revised: To construct this dataset, we ﬁrst deﬁne correspondences between the classes of CIFAR-10 and

Original: MNIST.
Revised: MNIST.

Original: 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 .
Revised: 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 .

Original: 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 .
Revised: 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 .

Original: It is to be noted that for both CIFAR-10 and MNIST, labels are always correlated with the respective images.
Revised: It is to be noted that for both CIFAR-10 and MNIST, labels are always correlated with the respective images.

Original: 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.
Revised: 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.

Original: This dataset represents the toughest setting of OOD generalization, where there is no differentiation between important features and spurious
Revised: This dataset represents the toughest setting of OOD generalization, where there is no differentiation between important features and spurious

Original: [Table]
Revised: [Table]

Original: correlations.
Revised: correlations.

Original: 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.
Revised: 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.

Original: Here the model could either rely on shape or color for classiﬁcation.
Revised: Here the model could either rely on shape or color for classiﬁcation.

Original: 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 [Table]
Revised: 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.

Original: 5. T RAINING AND E VALUATION S ETTINGS
Revised: 5. T RAINING AND E VALUATION S ETTINGS

Original: We consider two separate ResNet-18 (He et al., 2016) feature extractors for CIFAR-10 and MNIST respectively.
Revised: We consider two separate ResNet-18 (He et al., 2016) feature extractors for CIFAR-10 and MNIST respectively.

Original: 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.
Revised: 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.

Original: This architecture allows the computation of accuracy based on either a combination of both
Revised: This architecture allows the computation of accuracy based on either a combination of both

Original: CIFAR-10 and MNIST features, or based on features of only one of the datasets.
Revised: CIFAR-10 and MNIST features, or based on features of only one of the datasets.

Original: 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.
Revised: 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.

Original: We refer to this as the CIFAR-AvgMNIST dataset, while the corresponding one for MNIST is refered to as the MNIST-AvgCIFAR dataset.
Revised: We refer to this as the CIFAR-AvgMNIST dataset, while the corresponding one for MNIST is refered to as the MNIST-AvgCIFAR dataset.

Original: Similar to the work by
Revised: Similar to the work by

Original: Shah et al.
Revised: Shah et al.

Original: (2020), 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.
Revised: (2020), 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.

Original: 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.
Revised: 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.

Original: 5. E XPERIMENTAL R ESULTS IN VARIOUS T RAINING R EGIMES
Revised: 5. E XPERIMENTAL R ESULTS IN VARIOUS T RAINING R EGIMES

Original: We present the results of training on the MNIST-CIFAR-10 dataset using different algorithms in
Revised: We present the results of training on the MNIST-CIFAR-10 dataset using different algorithms in

Original: Table 2.
Revised: Table 2.

Original: The mean and standard deviation across ﬁve runs have been reported for each case.
Revised: The mean and standard deviation across ﬁve runs have been reported for each case.

Original: ERM Training:
Revised: ERM Training:

Original: By training a randomly initialized model on the MNIST-CIFAR-10 dataset using the cross-entropy loss (E1), we obtain an accuracy of 99 .
Revised: By training a randomly initialized model on the MNIST-CIFAR-10 dataset using the cross-entropy loss (E1), we obtain an accuracy of 99 .

Original: 84% on its corresponding test split.
Revised: 84% on its corresponding test split.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: To understand this, we reinitialize the linear classiﬁcation head randomly, and train the same using CIFAR-RandMNIST (E5) and
Revised: To understand this, we reinitialize the linear classiﬁcation head randomly, and train the same using CIFAR-RandMNIST (E5) and

Original: MNIST-RandCIFAR datasets (E6) respectively.
Revised: MNIST-RandCIFAR datasets (E6) respectively.

Original: 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.
Revised: 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.

Original: The MNIST-avgCIFAR accuracy of the latter case is high as expected.
Revised: The MNIST-avgCIFAR accuracy of the latter case is high as expected.

Original: 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 .
Revised: 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 .

Original: 53% ) and E3 ( 99 .
Revised: 53% ) and E3 ( 99 .

Original: 68% ) respectively.
Revised: 68% ) respectively.

Original: 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.
Revised: 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.

Original: We ﬁrstly explore the possibility of bridging the difference in the CIFARavgMNIST accuracy between E1 and E5 by merely retraining the linear layer.
Revised: We ﬁrstly explore the possibility of bridging the difference in the CIFARavgMNIST accuracy between E1 and E5 by merely retraining the linear layer.

Original: By reinitializing and naively retraining the linear layer with Cross-entropy loss, the accuracy on CIFAR-avgMNIST improves by less than 1% (E4).
Revised: By reinitializing and naively retraining the linear layer with Cross-entropy loss, the accuracy on CIFAR-avgMNIST improves by less than 1% (E4).

Original: 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.
Revised: 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.

Original: We penalize the (cid:96) ∞ norm of difference in original features and their reconstruction in addition to the minimization of cross-entropy loss.
Revised: We penalize the (cid:96) ∞ norm of difference in original features and their reconstruction in addition to the minimization of cross-entropy loss.

Original: The reconstruction based regularizer enforces the network to utilize both CIFAR and MNIST features for classiﬁcation.
Revised: The reconstruction based regularizer enforces the network to utilize both CIFAR and MNIST features for classiﬁcation.

Original: 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 −
Revised: 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 −

Original: I || F (E7).
Revised: I || F (E7).

Original: 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.
Revised: 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.

Original: Contrary to this, the proposed
Revised: Contrary to this, the proposed

Original: Feature Reconstruction Regularizer allows more ﬂexibility by imposing this constraint only on the domain of features seen during training.
Revised: Feature Reconstruction Regularizer allows more ﬂexibility by imposing this constraint only on the domain of features seen during training.

Original: This accounts for the simple feature replication as well, enabling to view the logit layer as an information bottleneck in the reconstruction.
Revised: This accounts for the simple feature replication as well, enabling to view the logit layer as an information bottleneck in the reconstruction.

Original: Finetuning (FT) and Fixed Linear Finetuning (FLFT):
Revised: Finetuning (FT) and Fixed Linear Finetuning (FLFT):

Original: 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.
Revised: 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.

Original: By ﬁnetuning an ERM trained base model using either of the two strategies (E9 and E10), we observe gains of less than 1% .
Revised: By ﬁnetuning an ERM trained base model using either of the two strategies (E9 and E10), we observe gains of less than 1% .

Original: We observe similar gains even by ﬁnetuning the full network with FRR (E11).
Revised: We observe similar gains even by ﬁnetuning the full network with FRR (E11).

Original: Contrary to this, by using FRR-FLFT even on the ERM trained network (E12), we obtain 7 .
Revised: Contrary to this, by using FRR-FLFT even on the ERM trained network (E12), we obtain 7 .

Original: 29% improvement over the base model.
Revised: 29% improvement over the base model.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Combining FRR-L and FRR-FLFT: While we obtain similar order of gains ( ∼ 7% ) using both
Revised: Combining FRR-L and FRR-FLFT: While we obtain similar order of gains ( ∼ 7% ) using both

Original: FRR-L and FRR-FLFT
Revised: FRR-L and FRR-FLFT

Original: 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.
Revised: 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.

Original: We therefore propose a training strategy that combines both FRR-L and FRR-FLFT.
Revised: We therefore propose a training strategy that combines both FRR-L and FRR-FLFT.

Original: Using this, we obtain gains of 16 .
Revised: Using this, we obtain gains of 16 .

Original: 2% over the ERM baseline as shown in E16, indicating that the combination of
Revised: 2% over the ERM baseline as shown in E16, indicating that the combination of

Original: 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.
Revised: 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.

Original: Although FRR-L followed by FRR-FT (E15) is also effective, it has about 6% lesser gains when compared to the proposed approach of
Revised: Although FRR-L followed by FRR-FT (E15) is also effective, it has about 6% lesser gains when compared to the proposed approach of

Original: FRR-L + FRR-FLFT.
Revised: FRR-L + FRR-FLFT.

Original: We note that following up FRR-L with ERM-FT (E13) or ERM-FLFT
Revised: We note that following up FRR-L with ERM-FT (E13) or ERM-FLFT

Original: (E14) also reﬁnes the learned features, making them more suitable for the classiﬁcation task, yielding 2 .
Revised: (E14) also reﬁnes the learned features, making them more suitable for the classiﬁcation task, yielding 2 .

Original: 6% and 4 .
Revised: 6% and 4 .

Original: 6% gains respectively over FRR-L.
Revised: 6% gains respectively over FRR-L.

Original: 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.
Revised: 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.

Original: We observe considerable gains of around 15% on MNIST-CIFAR-10 accuracy using CIFAR-RandMNIST training when compared to
Revised: We observe considerable gains of around 15% on MNIST-CIFAR-10 accuracy using CIFAR-RandMNIST training when compared to

Original: 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.
Revised: 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.

Original: 5. OOD G ENERALIZATION IN A R EAL W ORLD S ETTING
Revised: 5. OOD G ENERALIZATION IN A R EAL W ORLD S ETTING

Original: We show the efﬁcacy of FRR towards improving OOD generalization on the DomainBed (Gulrajani & Lopez-Paz, 2020) benchmark.
Revised: We show the efﬁcacy of FRR towards improving OOD generalization on the DomainBed (Gulrajani & Lopez-Paz, 2020) benchmark.

Original: We use the performance of the model on in-domain validation data (i.e. the in-domain strategy by Gulrajani & Lopez-Paz (2020)) to select the best hyper-parameters, and report the average performance and standard deviation across 5 random seeds.
Revised: We use the performance of the model on in-domain validation data (i.e. the in-domain strategy by Gulrajani & Lopez-Paz (2020)) to select the best hyper-parameters, and report the average performance and standard deviation across 5 random seeds.

Original: Baselines : We compare our method against standard ERM training, which has proven to be a frustratingly difﬁcult baseline (Gulrajani & Lopez-Paz, 2020), and also against several state of the art methods on this benchmark - SWAD (Cha et al., 2021), MIRO (Cha et al., 2022) and SMA (Arpit[Table] 2 ±
Revised: Baselines : We compare our method against standard ERM training, which has proven to be a frustratingly difﬁcult baseline (Gulrajani & Lopez-Paz, 2020), and also against several state of the art methods on this benchmark - SWAD (Cha et al., 2021), MIRO (Cha et al., 2022) and SMA (Arpit

Original: [Table]
Revised: [Table]

Original: et al., 2021).
Revised: et al., 2021).

Original: Finally, we show that our approach can be effectively integrated with stochastic weight averaging to obtain further gains.
Revised: Finally, we show that our approach can be effectively integrated with stochastic weight averaging to obtain further gains.

Original: See Appendix G for further experimental details.
Revised: See Appendix G for further experimental details.

Original: Main Results: The main results of our algorithm are reported in Table 3.
Revised: Main Results: The main results of our algorithm are reported in Table 3.

Original: 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.
Revised: 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.

Original: Further, our method obtains substantial gains of more than 3% over ERM across datasets.
Revised: Further, our method obtains substantial gains of more than 3% over ERM across datasets.

Original: The gains are especially pronounced for the larger datasets including DomainNet and TerraIncognita (8% and 5% resp.), indicating the efﬁcacy and scalability of our method.
Revised: The gains are especially pronounced for the larger datasets including DomainNet and TerraIncognita (8% and 5% resp.), indicating the efﬁcacy and scalability of our method.

Original: Further, it is clear from Table 3 that ﬁnetuning the feature extractor once the linear layer is ﬁxed provides a boost of over 1% on average over FRR-L.
Revised: Further, it is clear from Table 3 that ﬁnetuning the feature extractor once the linear layer is ﬁxed provides a boost of over 1% on average over FRR-L.

Original: This empirically validates our ﬁnetuning paradigm which we denote as ERM+FRR.
Revised: This empirically validates our ﬁnetuning paradigm which we denote as ERM+FRR.

Original: Finally, using our method in tandem with SWAD helps us achieve a new state-of-the-art on the
Revised: Finally, using our method in tandem with SWAD helps us achieve a new state-of-the-art on the

Original: DomainBed benchmark, outperforming other methods on three datasets while achieving comparable performance on two, and being better than existing SOTA by close to 1% on average.
Revised: DomainBed benchmark, outperforming other methods on three datasets while achieving comparable performance on two, and being better than existing SOTA by close to 1% on average.

Original: We report detailed results and further ablations in Appendices H and I.
Revised: We report detailed results and further ablations in Appendices H and J.

Original:  C ONCLUSION AND D ISCUSSION
Revised:  C ONCLUSION AND D ISCUSSION

Original: In this work, we consider the problem of OOD generalization through the lens of mitigating Simplicity Bias in Neural Network training.
Revised: In this work, we consider the problem of OOD generalization through the lens of mitigating Simplicity Bias in Neural Network training.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: We justify the proposed regularizer both theoretically and empirically on synthetic and semisynthetic datasets, and demonstrate its effectiveness in a real world OOD generalization setting.
Revised: We justify the proposed regularizer both theoretically and empirically on synthetic and semisynthetic datasets, and demonstrate its effectiveness in a real world OOD generalization setting.

Original: 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.
Revised: 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.

Original: We believe the proposed regularizer can potentially work effectively in several other settings that involve the use of linear layer training/ﬁnetuning, such as domain adaptation and transfer learning.
Revised: We believe the proposed regularizer can potentially work effectively in several other settings that involve the use of linear layer training/ﬁnetuning, such as domain adaptation and transfer learning.

Original: While the regularizer works effectively in a scenario where the network is ﬁrst trained using an algorithm such as ERM to learn features that are relevant to the task at hand, the robustness of the proposed algorithm in the presence of severely non-relevant features is yet to be explored.
Revised: While the regularizer works effectively in a scenario where the network is ﬁrst trained using an algorithm such as ERM to learn features that are relevant to the task at hand, the robustness of the proposed algorithm in the presence of severely non-relevant features is yet to be explored.

Original: A P ROOFS OF THEORETICAL RESULTS
Revised: A P ROOFS OF THEORETICAL RESULTS

Original: Proof of claim 3.1 :
Revised: In this section, we present a generalization of Claim 3.1 and Proposition 3.2 under weaker assumptions on the featurizer, and use them to prove the claims made in Sec 3. A more general setting Consider the dataset distribution D as deﬁned in eq equation 2.

Original: Proof.
Revised: Now, let f Θ ( x ) = Θ x , where Θ ∈ d × 2 . Further, ∀ i θ i, 1 + θ i, 2 = 1 . In simple words, the feature extractor maps the input to a d -dimensional feature representation, where each feature is a convex combination of the input feature variations.

Original: First, observe that w = 
Revised: Also note that θ =   11...

Original: [1 , 1] satisﬁes the constraint y · (cid:104) w, x (cid:105) ≥ 1 ∀ ( x, y
Revised: 10  corresponds to feature replication.

Original: ) ∈ Supp ( D ) . Now, consider any w satisfying this constraint.
Revised: We now rephrase the results from Sec 3.3 in this setting and provide proofs for the same.

Original: By considering ( y · (0 . 5 , 0 . 5) , y ) ∈ Supp ( D ) , we have that w 1 + w 2≥ 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 .
Revised: Claim A.1. (Restating Claim 3.1) - The max margin classiﬁer w minimizing (cid:107) w (cid:107) 2 and satisfying y (cid:104) w , Θ x (cid:105) ≥ 1 is given by w = (cid:2) 2 d , · · · , 2 d (cid:3) .

Original: The proof of the second part of the claim is identical to that of the ﬁrst part.
Revised: Proof.

Original: First, observe that ˜ w = [-(cid:2) 
-]
Revised: Consider the point x = [-(cid:2) 
-](0 . 5 , 0 . 5) .

Original: d +1 , · · · , 2 d +1 (cid:3) satisﬁes the constraint  y ·(cid:104) ˜ w, ˜ x (cid:105) ≥ 1 ∀ (˜ x, y
Revised: Then, due to the constraint on the max-margin classiﬁer, we have y (cid:104) w , Θ x (cid:105) ≥ 1 i.e.

Original: ) ∈ 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≥
Revised: [Equation]

Original: 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
Revised: The minimizer of (cid:96) 2 norm under this constraint would be when all w [-d +1≥ 1 , (cid:2) 
 d +1 , · · · ,-]j = 2 d for all j . Note that the ”effective classiﬁer” ˜ w in the input space in this case is (cid:104) w, Θ (cid:105) , i.e. the slope of the classiﬁer in the input space

Original: d +1 on
Revised: [Equation] =

Original: ˜ x ∈ R d +1 is equivalent to the classiﬁer w proj =(cid:2)
Revised: [Equation] dj =

Original: 2 d d +1 , 2 d +1 (cid:3)
Revised: 1 θ j, 1 .

Original: ∈ R 2 on x ∈ R 2: (cid:104) w proj , x (cid:105) = 
Revised: In particular, for the case of feature replication, (cid:80) dj =1 θ j, 2 = d − 1 and (cid:80) dj

Original: (cid:104)
Revised: =1 θ j, 1 = 1 , leading to a skewed classiﬁer. Now we restate and show the robustness of FRR in this setting. Proposition A.2. (Restating Prop 3.2) -

Original: ˜ w MM , ˜ x (cid:105) ∀ x ∈ R
Revised: Denote the feature reconstruction regularizer for this setting as -

Original: [Equation]
Revised: [Equation]

Original: Proof of Proposition 3.
Revised: Let w FRR be the minimizer of FRR ( w ) satisfying y (cid:104) w FRR

Original: Proof.
Revised: [Equation]

Original: Given any ˜ w, we ﬁrst compute L FRR ( ˜ w ) :
Revised: Θ x (cid:105) ≥ 0 (i.e. it is a perfect classiﬁer). Then, w FRR satisﬁes

Original: = min ˜
Revised: [Equation]

Original: φ L FRR ( ˜ w, ˜ φ ) .

Original: To do so, we note that the minimizing ˜ φ for a given ˜ w , denoted by ˜ φ ∗ ( ˜ w ) is given by:
Revised: where a = arg max i θ i, 2 , b = arg max i θ i, 1 . Proof. Consider the FRR for this dataset -

Original: [Equation]
Revised: [Equation]

Original: So, ˜ φ ∗ i ( ˜ w ) = E [ (cid:104) ˜ w, ˜ x (cid:105) ˜
Revised: For each i , by computing the minimum over u i by considering the quadratic in u i , we obtain [Equation]

Original: x i ] E [ (cid:104) ˜ w, ˜ x (cid:105) 2 ] . So, we have that:
Revised: We now consider each term in this expression

Original: [Equation]
Revised: [Equation]

Original: 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:
Revised: This is because E [ x 1 x 2 ] = 1 and E [ x 21 ] = E [ x 22 ] = 1 + 112 .

Original: E
Revised: E

Original: [Equation]
Revised: [Equation]

Revised: Finally,

Original: [Equation]
Revised: [Equation]

Revised: Putting it together,

Original: (cid:26)
Revised: [Equation] Let (cid:104) w , Θ [Equation] [Equation]

Revised: w , Θ ·

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Revised: +

Revised: [Equation]

Revised: w , Θ

Revised: [Equation]

Revised: [Equation]

Revised: = α .

Revised: Further, Let a, b be such that a = arg max i ( α − θ i, 1 ) 2 and b = arg max i (1 − α − θ i, 2 ) 2 .

Revised: Then,

Revised: [Equation]

Revised: To minimize the above expression w.r.t.

Revised: α , we compute the derivative of the above expression for each component of the max function

Revised: [Equation]

Revised: and

Revised: [Equation]

Revised: For α ≥ 1+ θ a, 1 − θ b, 2, the second term is greater and the derivative is positive.

Revised: Similarly, for α ≤ 1+ θ a, 1 − θ b, 2, the ﬁrst term is greater and the derivative is positive.

Original: if 1  ≤ i ≤ d ˜ w{+1 or 0 . 
 ≤ θ a, 1 , θ b, 2 ≤ 0 .+}
Revised: Hence, the minima is obtained at α = 1+ θ a, 1 − θ b, 2. Assuming 0 . 5 ≤ θ a, 1 , θ b, 2 ≤ {+1 or 0 . 
 ≤ θ a, 1 , θ b, 2 ≤ 0 .+}

Original:  [Equation] 
Revised: Now, in order to compute a, b , we can look at the maximization problem again - [Equation] Since θ i, 1 , θ i, 2 are bounded between [0 , 1] the function to maximize is monotonically decreasing in θ · , 1 , θ · , 2 . Hence, a = arg min i

Original: if i = d + 1 , where we used the fact that E [ n 21 ] = E [ n 22 ] = 1 / 12 .
Revised: θ i, 1 and b = arg min i θ i, 2 .

Original: Consequently, we have that L FRR ( ˜ w ) = 
Revised: Conversely, a = arg max i θ i, 2 , b = arg max

Original: 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) .
Revised: i θ i, 1 and α = 1+( θ b, 1 − θ a, 2 ).

Original: Consequently, any minimizer ˜ w ∗ of L FRR (
Revised: Hence,

Original: [Table]
Revised: [Equation]

Original:  w ) satisﬁes
Revised: Assuming that the maximum correlation of both the features is close, FRR will lead to a solution which gives roughly equal weights to both the features. Note that for the case of feature replication, θ b, 1 = 1 and θ a, 2 = 1 . Hence, (cid:104) w , Θ · , 2 (cid:105) (cid:104) w , Θ · , 1 (cid:105)

Original: [Equation]
Revised: = 1 .

Original: B J USTIFICATION OF F EATURE R EPLICATION H YPOTHESIS (FRH)
Revised: B J USTIFICATION OF F EATURE R EPLICATION H YPOTHESIS (FRH)

Original: In a practical scenario where features are not disentangled, our hypothesis translates to the following:
Revised: In a practical scenario where features are not disentangled, our hypothesis translates to the following:

Original: Conjecture: Simpler features of the input are represented more in the feature space of neural networks, while complex (hard-to-learn) features are sparse.
Revised: Conjecture: Simpler features of the input are represented more in the feature space of neural networks, while complex (hard-to-learn) features are sparse.

Original: Assumptions:
Revised: Assumptions:

Original: • We consider simple features such as background to be spurious, and complex features such as shape to be robust.
Revised: • We consider simple features such as background to be spurious, and complex features such as shape to be robust.

Original: • We consider an overparameterized network that has the capacity to learn more features than what exist, resulting in feature repetition.
Revised: • We consider an overparameterized network that has the capacity to learn more features than what exist, resulting in feature repetition.

Original: Justiﬁcation: We justify the conjecture by showing that all other possibilities discussed below cannot be true
Revised: Justiﬁcation: We justify the conjecture by showing that all other possibilities discussed below cannot be true.

Original: Assumption: DNNs learn only Simple Features Contradiction: Prior works (Rosenfeld et al., 2022; Kirichenko et al., 2022b) show that features learned by ERM are diverse, and last layer training on target domain is good enough to obtain robustness to spurious features.
Revised: Assumption: DNNs learn only Simple Features Contradiction: Prior works (Rosenfeld et al., 2022; Kirichenko et al., 2022b) show that features learned by ERM are diverse, and last layer training on target domain is good enough to obtain robustness to spurious features.

Original: This cannot be possible if the network has learned only spurious features.
Revised: This cannot be possible if the network has learned only spurious features.

Original: Assumption: DNNs learn only Complex Features Contradiction: The dominance of Simple features in the learning of DNNs is shown by Shah et al.
Revised: Assumption: DNNs learn only Complex Features Contradiction: The dominance of Simple features in the learning of DNNs is shown by Shah et al.

Original: Moreover, the existence of texture-bias (Geirhos et al., 2018) and background-bias (Xiao et al., 2020) have been demonstrated in prior works, which show the dominance of Simple features.
Revised: Moreover, the existence of texture-bias (Geirhos et al., 2018) and background-bias (Xiao et al., 2020) have been demonstrated in prior works, which show the dominance of Simple features.

Original: Assumption: DNNs learn a uniform distribution of both Simple and Complex Features.
Revised: Assumption: DNNs learn a uniform distribution of both Simple and Complex Features.

Original: Contradiction: SGD converges to an SVM solution due to its implicit bias (Soudry et al., 2018).
Revised: Contradiction: SGD converges to an SVM solution due to its implicit bias (Soudry et al., 2018).

Original: From Claim-3.1 (1), in the presence of balanced features that are correlated with the label, SVM solution gives equal weight to all features to maximize margin.
Revised: From Claim-3.1 (1), in the presence of balanced features that are correlated with the label, SVM solution gives equal weight to all features to maximize margin.

Original: This contradicts the existence of Simplicity Bias (Shah et al., 2020).
Revised: This contradicts the existence of Simplicity Bias (Shah et al., 2020).

Original: Assumption: DNNs learn more Complex Features and less Simple Features.
Revised: Assumption: DNNs learn more Complex Features and less Simple Features.

Original: Contradiction: Since Complex features are indeed more robust and are better correlated with the labels, the classiﬁer would rely more on these features.
Revised: Contradiction: Since Complex features are indeed more robust and are better correlated with the labels, the classiﬁer would rely more on these features.

Original: This contradicts the existence of Simplicity Bias (Shah et al., 2020).
Revised: This contradicts the existence of Simplicity Bias (Shah et al., 2020).

Original: Therefore, the only feasible option which supports the empirical observations in the prior works discussed above is that DNNs learn more Simple Features and Complex features are sparse, which justiﬁes our conjecture.
Revised: Therefore, the only feasible option which supports the empirical observations in the prior works discussed above is that DNNs learn more Simple Features and Complex features are sparse, which justiﬁes our conjecture.

Original: C E MPIRICAL E VIDENCE FOR F EATURE R EPLICATION H YPOTHESIS (FRH) C. S YNTHETIC D ATASETS
Revised: C E MPIRICAL E VIDENCE FOR F EATURE R EPLICATION H YPOTHESIS (FRH) C. S YNTHETIC D ATASETS

Original: We present empirical validation to support the Feature Replication Hypothesis (FRH) on several semi-real datasets and describe them in detail below:
Revised: We present empirical validation to support the Feature Replication Hypothesis (FRH) on several semi-real datasets and describe them in detail below:

Original: Coloured-MNIST-2 - In this dataset, we use images of digits superimposed on either of the two colours- red or green.
Revised: Coloured-MNIST-2 - In this dataset, we use images of digits superimposed on either of the two colours- red or green.

Original: The difference from Coloured-MNIST is that we consider only two colours for the background, rather than a range.
Revised: The difference from Coloured-MNIST is that we consider only two colours for the background, rather than a range.

Original: We notice extreme simplicity bias in this case, with the network learning 32 colour features and 0 shape features.
Revised: We notice extreme simplicity bias in this case, with the network learning 32 colour features and 0 shape features.

Original: Coloured-MNIST-MultiDigit -
Revised: Coloured-MNIST-MultiDigit -

Original: This is similar to the Coloured-MNIST dataset described in Section-3.1, with the exception that each of the classes is now composed of two digits.
Revised: This is similar to the Coloured-MNIST dataset described in Section-3.1, with the exception that each of the classes is now composed of two digits.

Original: More speciﬁcally, the digits ‘1’ and ‘7’ and mapped to Class 0 and digits ‘5’ and ‘8’ are mapped to Class 1.
Revised: More speciﬁcally, the digits ‘1’ and ‘7’ and mapped to Class 0 and digits ‘5’ and ‘8’ are mapped to Class 1.

Original: We note that ‘1’ and ‘5’ are chosen from the original Coloured-MNIST dataset, while the second digit (e.g.‘7’) in each class is selected to be one that is similar to the ﬁrst digit (‘1’) in the same class.
Revised: We note that ‘1’ and ‘5’ are chosen from the original Coloured-MNIST dataset, while the second digit (e.g.‘7’) in each class is selected to be one that is similar to the ﬁrst digit (‘1’) in the same class.

Original: This dataset is constructed speciﬁcally to show that the issue of Simplicity Bias and FRH exists even when there is higher variation in the shape feature, and is reported as ColouredMNIST-MultiDigit below.
Revised: This dataset is constructed speciﬁcally to show that the issue of Simplicity Bias and FRH exists even when there is higher variation in the shape feature, and is reported as ColouredMNIST-MultiDigit below.

Original: We see that while more shape features are learnt as compared to Coloured-MNIST, the network still relies more on colour to make its decisions.
Revised: We see that while more shape features are learnt as compared to Coloured-MNIST, the network still relies more on colour to make its decisions.

Original: Digit-Coloured-MNIST :
Revised: Digit-Coloured-MNIST :

Original: This is similar to the Coloured-MNIST dataset described in Section-3.1, with the exception that the digit is coloured rather than the background.
Revised: This is similar to the Coloured-MNIST dataset described in Section-3.1, with the exception that the digit is coloured rather than the background.

Original: This dataset is constructed speciﬁcally to show that the issue of SB and FRH exist even when the region that is coloured, which is the extent to which simple features exist in the image is much lesser, and is reported as DigitColouredMNIST below.
Revised: This dataset is constructed speciﬁcally to show that the issue of SB and FRH exist even when the region that is coloured, which is the extent to which simple features exist in the image is much lesser, and is reported as DigitColouredMNIST below.

Original: Although this dataset also demonstrates the presence of SB, we note that the average correlation of features with shape is higher when compared to the above datasets. 
Revised: Although this dataset also demonstrates the presence of SB, we note that the average correlation of features with shape is higher when compared to the above datasets. [Table]

Original: C. R EAL W ORLD EXAMPLE
Revised: C. R EAL W ORLD EXAMPLE

Original: We attempt to demonstrate feature replication in a model trained with ERM on the Real domain of
Revised: We attempt to demonstrate feature replication in a model trained with ERM on the Real domain of

Original: OfﬁceHome.
Revised: OfﬁceHome.

Original: We train a ResNet-50 on this domain, and perform PCA on the features learnt by this 
Revised: We train a ResNet-50 on this domain, and perform PCA on the features learnt by this network.

Original: [Table] network. The network learns 2048 features per example, and we compute the 2048 × 2048 sized covariance matrix of the features over samples from a test domain (Clipart).
Revised: The network learns 2048 features per example, and we compute the 2048 × 2048 sized covariance matrix of the features over samples from a test domain (Clipart).

Original: We then compute the eigenvalues of this matrix, and ﬁnd that 500 principal components can explain about 97.5% of the variance, i.e. the matrix is extremely low rank, as shown in Fig.3.
Revised: We then compute the eigenvalues of this matrix, and ﬁnd that 500 principal components can explain about 97.5% of the variance, i.e. the matrix is extremely low rank, as shown in Fig.3.

Original: This points to the fact that a lot of the learnt features are linearly dependent and highly correlated with each other.
Revised: This points to the fact that a lot of the learnt features are linearly dependent and highly correlated with each other.

Original: This trend is similar to what we observed on ColouredMNIST, where a large number of features were highly correlated with the colour, and in turn with each other (Fig.
Revised: This trend is similar to what we observed on ColouredMNIST, where a large number of features were highly correlated with the colour, and in turn with each other (Fig.

Original: 5 of the appendix).
Revised: 5 of the appendix).

Original: We note that in all the additional datasets considered, simpler features are represented more in the network while complex (hard-to-learn) features are sparse.
Revised: We note that in all the additional datasets considered, simpler features are represented more in the network while complex (hard-to-learn) features are sparse.

Original: This empirically justiﬁes our hypothesis in Section-3.
Revised: This empirically justiﬁes our hypothesis in Section-3.

Original: D S YNTHETIC DATASETS D. C OLOURED MNIST
Revised: D S YNTHETIC DATASETS D. C OLOURED MNIST

Original: In order to empirically demonstrate feature replication, we use a binarized version of the coloured
Revised: In order to empirically demonstrate feature replication, we use a binarized version of the coloured

Original: MNIST dataset (Gulrajani & Lopez-Paz, 2020).
Revised: MNIST dataset (Gulrajani & Lopez-Paz, 2020).

Original: To construct this dataset, we ﬁrstly assign two digits of the MNIST dataset, namely “1” and “5”, to classes 0 and 1 respectively.
Revised: To construct this dataset, we ﬁrstly assign two digits of the MNIST dataset, namely “1” and “5”, to classes 0 and 1 respectively.

Original: For the in-domain training distribution, we associate colours in the range R 0 =
Revised: For the in-domain training distribution, we associate colours in the range R 0 =

Original: [(115 , 0 , 0) − (256 , 141 , 0))]
Revised: [(115 , 0 , 0) − (256 , 141 , 0))]

Original: (i.e. red) to label 

Revised: (i.e. red) to label

Original: (i.e. the digit “1”) and the range R 1 =
Revised: (i.e. the digit “1”) and the range R 1 =

Original: [(0 , 115 , 0) − (141 , 256 , 0)]
Revised: [(0 , 115 , 0) − (141 , 256 , 0)]

Original: (i.e. green) to the label 1 (i.e. the digit “5”), where colors are represented in the RGB space.
Revised: (i.e. green) to the label 1 (i.e. the digit “5”), where colors are represented in the RGB space.

Original: 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 .
Revised: 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 .

Original: 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.
Revised: 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.

Original: This overlap reduces the correlation of colour features with labels, while shape features have a correlation of 1 with the labels.
Revised: This overlap reduces the correlation of colour features with labels, while shape features have a correlation of 1 with the labels.

Original: In Figure 4, we show examples of images from the train and test distributions of this dataset.
Revised: In Figure 4, we show examples of images from the train and test distributions of this dataset.

Original: In Figure 5, we pictorially depict the correlations between the 32 features learnt by the network.
Revised: In Figure 5, we pictorially depict the correlations between the 32 features learnt by the network.

Original: We can see a block structure emerging, indicating that there is a high amount of feature replication.
Revised: We can see a block structure emerging, indicating that there is a high amount of feature replication.

Original: D. T OY DATASET
Revised: D. T OY DATASET

Original: In line with the theoretical formulation described in Section-3.3, 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.
Revised: In line with the theoretical formulation described in Section-3.3, we further empirically validate the brittleness of SVM models and the highlight the effectiveness of the proposed Feature Reconstruc- tion Regularizer in the presence of replicated features.

Original: We consider a linearly separable toy distribution consisting of two factors of variation as shown in Figure 1.
Revised: We consider a linearly separable toy distribution consisting of two factors of variation as shown in Figure 1.

Original: 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
Revised: 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

Original: Unif[-0.5,0.5] independently along each dimension to the respective means.
Revised: Unif[-0.5,0.5] independently along each dimension to the respective means.

Original: We sample an Out-ofDistribution (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.
Revised: We sample an Out-ofDistribution (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.

Original: Therefore, while the train distribution can be classiﬁed by considering either the features aligned with the X-coordinate or the
Revised: Therefore, while the train distribution can be classiﬁed by considering either the features aligned with the X-coordinate or the

Original: Y-coordinate, the test set performance crucially depends on the variation along X-coordinate alone.
Revised: Y-coordinate, the test set performance crucially depends on the variation along X-coordinate alone.

Original: 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.
Revised: 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.

Original: To select the best hyperparameter for both SVM and FRR, we consider the presence of a validation set whose distribution is similar to the test distribution.
Revised: To select the best hyperparameter for both SVM and FRR, we consider the presence of a validation set whose distribution is similar to the test distribution.

Original: As shown in Figure 1, we observe that the SVM model starts relying more on the replicated features alone in case of feature replication, compromising its performance on the
Revised: As shown in Figure 1, we observe that the SVM model starts relying more on the replicated features alone in case of feature replication, compromising its performance on the

Original: OOD data.
Revised: OOD data.

Original: 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.
Revised: 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.

Original: E A LGORITHM
Revised: E A LGORITHM

Original: Data: Training data D S
Revised: Our training procedure is detailed in Alg 1. F D ETAILS ON THE OOD G ENERALIZATION 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 (Pan & Yang, 2009; Ganin et al., 2016), and Domain

Revised: Generalization , where typically data from several source distributions is assumed to be available, and the target domain in unseen during training (Blanchard et al., 2011; Li et al., 2018a; Gulrajani & Lopez-Paz, 2020)).

Revised: 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

Revised: Data: Training data D

Original:  = { ( x i , y i ) : i ∈ [ n ] } , model ( θ, W ) , feature reconstruction model φ ,
Revised: [Table] = { ( x i , y i ) : i ∈ [ n ] } , model ( θ, W ) , feature reconstruction model φ ,

Original: [Table]
Revised: [Table]

Revised: FRR ,

Original: [Table]
Revised: [Table]

Original: * Standard trainingof model parameters θand W .
Revised: * Standard training model parameters θ W .

Original: */ 2 Freeze θ to be

Original: [Table]
Revised: [Table]

Original:  /* Initializing model for training with FRR.
Revised: Freeze θ to be θ std /* Initializing model for training with FRR.

Original: [Equation]
Revised: [Table]

Original: * FRR-L: Training W, φ with FRR defined in eqn.
Revised: [Table] [Table]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: , ( x

Original: [Table]
Revised: [Table]

Original: [Equation] [Table] [Equation]
Revised: , y i )) + λ

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: L FRR ( x i , θ std , W, φ )) .

Original: / * FRR-FLFT:
Revised: / * FRR FRR

Original: Finetuning θ with FRR according to eqn.
Revised: [Table] [Table]

Original: [Table]
Revised: [Table]

Original: F D ETAILS ON THE OOD G ENERALIZATION SETTING CONSIDERED
Revised: Finetuning according to

Original: 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 (Pan & Yang, 2009; Ganin et al., 2016), and Domain Generalization , where typically data from several source distributions is assumed to be available, and the target domain in unseen during training (Blanchard et al., 2011; Li et al., 2018a; Gulrajani & Lopez-Paz, 2020)).
Revised: Result:

Original: 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 (Muandet et al., 2013; Ganin et al., 2016; Li et al., 2018b; Arjovsky et al., 2019; Shi et al., 2021).
Revised: Trained model ( θ FLFT , W FRR ) . that cause a distribution shift in input data without change in their label distribution (Muandet et al., 2013; Ganin et al., 2016; Li et al., 2018b; Arjovsky et al., 2019; Shi et al., 2021).

Original: 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.
Revised: 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.

Original: We investigate this crucial setting which has been relatively less researched, and refer to it as Aggregated Domain Generalization , as introduced by Thomas et al.
Revised: We investigate this crucial setting which has been relatively less researched, and refer to it as Aggregated Domain Generalization , as introduced by Thomas et al.

Original: 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 (Hendrycks & Dietterich, 2019), 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.
Revised: 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 (Hendrycks & Dietterich, 2019), 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.

Original: While there have been several approaches to improve the performance of models in the setting of Domain Generalization, Gulrajani & Lopez-Paz (2020) 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) .
Revised: While there have been several approaches to improve the performance of models in the setting of Domain Generalization, Gulrajani & Lopez-Paz (2020) 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) .

Original: 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.
Revised: 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.

Original: G E XPERIMENTAL D ETAILS ON D OMAIN B ED
Revised: G E XPERIMENTAL D ETAILS ON D OMAIN B ED

Original: We test our approach on the DomainBed benchmark (Gulrajani & Lopez-Paz, 2020) comprising of ﬁve different datasets, each of which have k domains.
Revised: We test our approach on the DomainBed benchmark (Gulrajani & Lopez-Paz, 2020) comprising of ﬁve different datasets, each of which have k domains.

Original: For each dataset, we train a model on k −domains, and test it on the left out domain.
Revised: For each dataset, we train a model on k −domains, and test it on the left out domain.

Original: The average out-of-domain performance across the k held-out domains is then reported.
Revised: The average out-of-domain performance across the k held-out domains is then reported.

Original: In this section we describe the hyper-parameter selection strategy and the ranges considered for our approach.
Revised: In this section we describe the hyper-parameter selection strategy and the ranges considered for our approach.

Original: In line with the DomainBed testbench, we use ImageNet pretrained ResNet-50 models for all algorithms.
Revised: In line with the DomainBed testbench, we use ImageNet pretrained ResNet-50 models for all algorithms.

Original: We use random search to select hyperparameters for our algorithm, and use the suggested hyperparameters for the other baselines.
Revised: We use random search to select hyperparameters for our algorithm, and use the suggested hyperparameters for the other baselines.

Original: We train for(5000 for DomainNet) steps in the FRR-L phase, and 5000 (10000 for DomainNet) steps in the FRRFLFT phase.
Revised: We train for(5000 for DomainNet) steps in the FRR-L phase, and 5000 (10000 for DomainNet) steps in the FRRFLFT phase.

Original: The batch size is ﬁxed to 32, and SWAD hyper-parameters are the same as those used by Cha et al.
Revised: The batch size is ﬁxed to 32, and SWAD hyper-parameters are the same as those used by Cha et al.

Original: We use the in-domain accuracy protocol from Gulrajani & Lopez-Paz (2020) to select hyper-parameters for each domain of each dataset, and search over 8 random combinations of hyper-parameters for each.
Revised: We use the in-domain accuracy protocol from Gulrajani & Lopez-Paz (2020) to select hyper-parameters for each domain of each dataset, and search over 8 random combinations of hyper-parameters for each.

Original: The range of the hyperparameters is shown in Table 5.
Revised: The range of the hyperparameters is shown in Table 5.

Original: 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.
Revised: 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.

Original: All our experiments were done on single V100 GPUs.
Revised: All our experiments were done on single V100 GPUs.

Original: [Table]
Revised: [Table]

Original: H A BLATIONS ON D OMAIN B ED
Revised: H A BLATIONS ON D OMAIN B ED

Original: Comparing the choices for φ : In Table 6, we experiment with various architectures for the decoder φ when computing FRR according to equation 1.
Revised: Comparing the choices for φ : In Table 6, we experiment with various architectures for the decoder φ when computing FRR according to equation 1.

Original: We consider using a two layer neural network as the decoder φ (FRR-LDeeper), and also consider setting φ = W T (FRR-LShared), i.e. explicitly tying the weights of the decoder and the classiﬁer layer.
Revised: 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.

Original: Overall, both these variants are worse than the default single layer, free parameterization of φ .
Revised: Overall, both these variants are worse than the default single layer, free parameterization of φ .

Original: 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.
Revised: 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.

Original: Both these have a detrimental effect on the overall performance of the model.
Revised: Both these have a detrimental effect on the overall performance of the model.

Original: [Table]
Revised: [Table]

Original: Sensitivity Analysis : We vary λ FRR and plot out the OOD performance in Fig 6.
Revised: Sensitivity Analysis : We vary λ FRR and plot out the OOD performance in Fig 6.

Original: We ﬁnd that the performance is stable for a wide range of the hyper-parameter on most domains. 
Revised: We ﬁnd that the performance is stable for a wide range of the hyper-parameter on most domains. I

Revised: P SEUDO -C ODE FOR FRR

Revised: Below we provide the python code for FRR-L in the DomainBed framework.

Revised: c

Original: I D OMAIN WISE ACCURACIES
Revised: l a s s ERMWithFRR L( Algorithm ) : def i n i t ( s e l f , input shape , num classes , num domains , hparams ) : super (ERMWithFRR L , s e l f ) . i n i t ( input shape , num classes , num domains , hparams ) s e l f . f e a t u r i z e r = networks .

Revised: F e a t u r i z e r ( input shape , s e l f . hparams ) s e l f .

Revised: c

Revised: l a s s

Revised: i f

Revised: i e r = networks .

Revised: C l a s s

Revised: i f

Original: In this section, we show detailed results of Table 3 in the main text.
Revised: i e r ( s e l f . f e a t u r i z e r . n outputs , num classes , s e l f .

Revised: hparams [ ’ n o n l i n e a r c l a s s i

Revised: f i e r ’ ] , ) for params in s e l f .

Original: The numbers for the baselines are taken from Gulrajani & Lopez-Paz (2020), Cha et al.
Revised: f e a t u r i z e r . parameters ( ) : params . r e q u i r e s g r a d

Revised: =

Revised: False s e l f .

Original: (2021) and Arpit et al.
Revised: c l a s s i f i e r i n v = networks .

Revised: C l a s s

Revised: i f

Original: (2021), while the results for MIRO (Cha et al., 2022) were reproduced using their code-base.
Revised: i e r ( num classes , s e l f . f e a t u r i z e r . n outputs , s e l f . hparams [ ’ n o n l i n e a r c l a s s i

Revised: f i e r ’ ] ,

Original: − 6 − 5 − 4 − 3 − 2 − 11 log( λ FRR ) → 0 . 885 0 .0 . 
Revised: − 6 − 5 − 4 − 3 − 2 − 11 log( λ FRR ) → 0 . 885 0 .0 . [Table] [Table]

Original: OO D
Revised: OO D

Original: A cc u r a c y →
Revised: A cc u r a c y →

Original: Art − 6 − 5 − 4 − 3 − 2 − 1 0 log( λ FRR ) → 0 .0 . 
Revised: Art − 6 − 5 − 4 − 3 − 2 − 1 0 log( λ FRR ) → 0 .0 . [Table] [Table] [Table]

Revised: [Table]

Original: OO D
Revised: OO D

Original: A cc u r a c y →
Revised: A cc u r a c y →

Original: Cartoon − 6 − 5 − 4 − 3 − 2 − 1 0 log( λ FRR ) → 0.0 . 825 0 .0 .
Revised: Cartoon − 6 − 5 − 4 − 3 − 2 − 1 0 log( λ FRR ) → 0 .

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table] [Table]
Revised: Photo

Original: [Table]
Revised: [Table]

Original: OO
Revised: −

Original: [Table]
Revised: [Table]

Original: u r
Revised: ( λ FRR )

Original: [Table]
Revised: [Table]

Original: Photo − 6 −

Original: [Table]
Revised: [Table]

Original: −
Revised: [Table]

Original: [Table]
Revised: [Table]

Revised: −

Original: 3 − 2 − 1 0 log( λ FRR
Revised: [Table] [Table] [Table] OO

Original: [Table]
Revised: [Table]

Revised: A

Original: [Table]
Revised: [Table]

Revised: c

Original: [Table]
Revised: [Table]

Revised: − 6 − 5 − 4 − 3 − 2 − 1 0 log( λ FRR )

Original: [Table] 
Revised: [Table] [Table] [Table] [Table]

Revised: [Table]

Original: [Table]
Revised: [Table]

Original: Sketch
Revised: OO D A cc u r a c y

Original: [Table]
Revised: [Table]

Original:  − 6 − 5 − 4 − 3 − 
Revised: LabelMe − 6 − 5 − 4 − 3 − 2 −

Original: [Table]
Revised: [Table]

Original:  FRR) →
Revised: log( λ FRR

Original: [Table]
Revised: [Table]

Original: OO D
Revised: [Table] [Table] [Table]

Original: A cc u r a c y → CalTech − 6 − 5 − 4 − 3 − 2 − 1 0 log( λ FRR )
Revised: [Table]

Original: → 0 .0 .
Revised: [Table]

Original: OO D
Revised: [Table]

Original: A cc u r a c y → LabelMe − 6 − 5 − 4 − 3 − 2 − 1 0 log( λ FRR ) → 0 .
Revised: [Table]

Original: [Table]
Revised: [Table]

Original: OOD A
Revised: OO

Original: [Table]
Revised: [Table]

Original: u r a c y
Revised: r a c y

Original: [Table]
Revised: [Table]

Original: PascalVoc − 6 − 5 − 4 − 3 − 2 − 1 0 log( λ FRR ) →
Revised: PascalVoc − 6 − 5 − [Table] − 3 − 2 − 1 0 log( λ FRR ) [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Original: [Table]
Revised: [Table]

Original: OO D
Revised: OO D

Original: A cc u r a c y →
Revised: A cc u r a c y →

Original: SUN (b) VLCS − 6 − 5 − 4 − 3 − 2 − 1 0 log( λ FRR ) → 0 .0 .
Revised: SUN (b) VLCS − 6 − 5 − 4 − 3 − 2 − 1 0 log( λ FRR ) → 0 .0 .

Original: OO D
Revised: OO D

Original: A cc u r a c y →
Revised: A cc u r a c y →

Original: Art − 6 − 5 − 4 − 3 − 2 − 1 0 log( λ FRR ) → 0 . 30 0 .0 .
Revised: Art − 6 − 5 − 4 − 3 − 2 − 1 0 log( λ FRR ) → 0 . 30 0 .0 .

Original: OO D
Revised: OO D

Original: A cc u r a c y →
Revised: A cc u r a c y →

Original: Clipart − 6 − 5 − 4 − 3 − 2 − 1 0 log( λ FRR )
Revised: Clipart − 6 − 5 − 4 − 3 − 2 − 1 0 log( λ FRR )

Original: → 0 .0 .0 .0 .0 .
Revised: → 0 .0 .0 .0 .0 .

Original: OO D
Revised: OO D

Original: A cc u r a c y →
Revised: A cc u r a c y →

Original: Product − 6 − 5 − 4 − 3 − 2 − 1 0 log( λ FRR )
Revised: Product − 6 − 5 − 4 − 3 − 2 − 1 0 log( λ FRR )

Original: → 0 .0 .0 .0 .0 .0 .
Revised: → 0 .0 .0 .0 .0 .0 .

Original: OO D
Revised: OO D

Original: A cc u r a c y →
Revised: A cc u r a c y →

Original: Real (
Revised: Real (c) OfﬁceHome

Original: [Figure]
Revised: i f ’ e x a c t r e c o n s t r u c t i o n ’ in s e l f .

Original: OfﬁceHome
Revised: hparams and s e l f . hparams [ ’ e x a c t r e c o n s t r u c t i o n ’ ]

Original: [Table]
Revised: e l s e True , ) s = e q e t i a l ( s e l f . f e a t u r i z

Revised: e

Revised: , s e l c l s s i

Revised: f i e r )

Revised: s e l o p

Revised: i m i z e r =

Revised: t o r c optim .

Revised: Adam( ( l i s t ( s e l f .

Revised: network .

Revised: parameters ( ) )

Revised: + l i s t ( s e l f .

Revised: c

Revised: l a s s

Revised: i f

Revised: i e r i n v .

Revised: parameters ( ) ) ) , l r = s e l f . hparams [ ’ l r ’ ] , weight decay =

Revised: s e l f .

Revised: hparams [ ’ weight decay ’ ] , ) s e l f .

Revised: r e c o n s t r u c t i o n w t =

Revised: s e l f .

Revised: hparams [ ’ r e c o n s t r u c t

Revised: i

Revised: o n w t ’ ]

Revised: s e l f .

Revised: norm =

Revised: f l o a t ( s e l f .

Revised: hparams [ ’norm ’ ] ) def update ( s e l f , minibatches , u n l a b e l e d =None ) : a l l x = t

Revised: o r c h .

Revised: c a t ( [ x for x , y in minibatches ] ) a l l y =

Revised: t

Revised: o r c h .

Revised: c a t ( [ y for x , y in minibatches ] ) pred , rec , f e a t =

Revised: s e l f .

Revised: g e t f e a t s a n d r e c ( a l l x ) l o s s = F .

Revised: c

Revised: r o s s e n t r o p y ( pred , a l l y )

Revised: r e c o n s t r u c t

Revised: i

Revised: o

Revised: n l o s s = ( t o r c h .

Revised: sum ( t o r c h .

Revised: max

Revised: ( t o r c h .

Revised: abs ( f e a t − rec ) , dim = 1 ) [ 0 ] ) / a l l x . shape [ 0 ] )

Revised: lo s s =

Revised: l o s s + s e l f .

Revised: r e c o n s t r u c t i o n w t ∗ r e c o n s t r u c t i o n l o s s s e l f .

Revised: o p t i m i

Revised: z

Revised: e r .

Revised: z e r o g r a d ( ) l

Revised: o s s .

Revised: backward ( ) s e l f .

Revised: o p t i m i

Revised: z

Revised: e r .

Revised: s t e p ( ) return { ’ l o s s ’ :

Revised: l o s s . item ( ) , ’ r e c o n s t r u c t i o n l o s s ’ :

Revised: r e c o n s t r u c t i o n l o s s .

Revised: item ( ) , } def p r e d i c t ( s e l f , x ) : return s e l f .

Revised: network ( x ) def g e t f e a t s a n d r e c ( s e l f , x ) : f e a t s = s e l f .

Revised: network [ 0 ] ( x ) pred =

Revised: s e l f .

Revised: network [ 1 ] ( f e a t s ) rec =

Revised: s e l f .

Revised: c

Revised: l a s s

Original: [Table]
Revised: i f i e r i n v ( pred ) return pred , rec , f e a t s J D OMAIN WISE ACCURACIES In this section, we show detailed results of Table 3 in the main text.

Original: [Table]
Revised: The numbers for the baselines are taken from Gulrajani & Lopez-Paz (2020), Cha et al. (2021) and Arpit et al. (2021), while the results for MIRO (Cha et al., 2022) were reproduced using their code-base.

Original: [Table]
Revised: [Table]

Original: [Table] 
Revised: [Table] [Table] [Table] [Table]