Original: MaskTune : Mitigating Spurious Correlations by Forcing to Explore
Revised: MaskTune : Mitigating Spurious Correlations by Forcing to Explore

Original: Abstract
Revised: Abstract

Original: A fundamental challenge of over-parameterized deep learning models is learningmeaningful data representations that yield good performance on a downstreamtask without over-ﬁtting spurious input features.
Revised: A fundamental challenge of over-parameterized deep learning models is learning meaningful data representations that yield good performance on a downstream task without over-fitting spurious input features.

Original: This work proposes MaskTune ,a masking strategy that prevents over-reliance on spurious (or a limited numberof) features.
Revised: This work proposes MaskTune , a masking strategy that prevents over-reliance on spurious (or a limited number of) features.

Original: MaskTune forces the trained model to explore new features duringa single epoch ﬁnetuning by masking previously discovered features.
Revised: MaskTune forces the trained model to explore new features during a single epoch finetuning by masking previously discovered features.

Original: MaskTune ,unlike earlier approaches for mitigating shortcut learning, does not require anysupervision, such as annotating spurious features or labels for subgroup samplesin a dataset.
Revised: MaskTune , unlike earlier approaches for mitigating shortcut learning, does not require any supervision, such as annotating spurious features or labels for subgroup samples in a dataset.

Original: Our empirical results on biased MNIST, CelebA, Waterbirds, and
Revised: Our empirical results on biased MNIST, CelebA, Waterbirds, and

Original: ImagenNet-9L datasets show that MaskTune is effective on tasks that often sufferfrom the existence of spurious correlations.
Revised: ImagenNet-9L datasets show that MaskTune is effective on tasks that often suffer from the existence of spurious correlations.

Original: Finally, we show that MaskTuneoutperforms or achieves similar performance to the competing methods whenapplied to the selective classiﬁcation task. 
Revised: Finally, we show that MaskTune outperforms or achieves similar performance to the competing methods when applied to the selective classification (classification with rejection option) task. Code for MaskTune is available at https://github.com/aliasgharkhani/ Masktune .

Original:  Introduction
Revised:  Introduction

Original: Spurious correlations are coincidental feature associations formed between a subset of the input andtarget variables, which may be caused by factors such as data selection bias (Torralba & Efros, 2011;
Revised: Spurious correlations are coincidental feature associations formed between a subset of the input and target variables, which may be caused by factors such as data selection bias [Torralba and Efros, 2011,

Original: Jabri et al., 2016).
Revised: Jabri et al., 2016].

Original: The presence of spurious correlations in training data can cause over-parameterizeddeep neural networks to fail, often drastically, when such correlations do not hold in test data (Sagawaet al., 2019) or when encountering domain shift (Arjovsky et al., 2019).
Revised: The presence of spurious correlations in training data can cause over-parameterized deep neural networks to fail, often drastically, when such correlations do not hold in test data [Sagawa et al., 2019] or when encountering domain shift [Arjovsky et al., 2019].

Original: Consider the classiﬁcationproblem of cows and camels (Beery et al., 2018), where most of the images of cows vs. camelsare captured on green ﬁelds vs. desert backgrounds due to selection bias (and perhaps the natureof the problem that camels are often in the desert).
Revised: Consider the classification problem of cows and camels [Beery et al., 2018], where most of the images of cows vs. camels are captured on green fields vs. desert backgrounds due to selection bias (and perhaps the nature of the problem that camels are often in the desert).

Original: A model trained on such data may rely on thebackground as the key discriminative feature between cows and camels, thus failing on images ofcows on non-green backgrounds or camels on non-desert backgrounds.
Revised: A model trained on such data may rely on the background as the key discriminative feature between cows and camels, thus failing on images of cows on non-green backgrounds or camels on non-desert backgrounds.

Original: In over-parameterized regimes, there are often several solutions with almost identical loss values, andthe optimizer (e.g., SGD) typically selects a low-capacity one (Wilson et al., 2017).
Revised: In over-parameterized regimes, there are often several solutions with almost identical loss values, and the optimizer (e.g., SGD) typically selects a low-capacity one [Wilson et al., 2017, Valle-Perez et al., 2018, Arpit et al., 2017, Kalimeris et al., 2019].

Original: In the presence ofspurious correlations, the optimizer chooses to leverage them as they generally demand less capacitythan the expected semantic cues of interest, e.g., relying on the local color or texture of grass insteadof the elaborate visual features that give a cow its appearance (Bruna & Mallat, 2013; Bruna et al. ,2015; Brendel & Bethge, 2019).
Revised: In the presence of spurious correlations, the optimizer might choose to leverage them as they generally demand less capacity than the expected semantic cues of interest, e.g., relying on the local color or texture of grass instead of the elaborate visual features that give a cow its appearance [Bruna and Mallat, 2013, Bruna et al., 2015, Brendel and Bethge, 2019, Khani and Liang, 2021]. [Figure]

Original: In previous work, a supervised loss function has been employed to reduce the effect of spuriouscorrelations (Sagawa et al., 2019).
Revised: In previous work, a supervised loss function has been employed to reduce the effect of spurious correlations [Sagawa et al., 2019].

Original: However, identifying and annotating the spurious correlationsin a large dataset as a training signal is impractical.
Revised: However, identifying and annotating the spurious correlations in a large dataset as a training signal is impractical.

Original: Other works have attempted to force modelsto discard context and background (as spurious features) through input morphing using perceptualsimilarities (Taghanaki et al., 2021) or learning casual variables (Javed et al., 2020).
Revised: Other works have attempted to force models to discard context and background (as spurious features) through input morphing using perceptual similarities [Taghanaki et al., 2021] or learning casual variables [Javed et al., 2020].

Original: Discardingcontext and background, however, is incompatible with the human visual system that relies on[Figure] contextual information when detecting and recognizing objects (Palmer, 1975; Biederman et al.
Revised: Discarding context and background, however, is incompatible with the human visual system that relies on contextual information when detecting and recognizing objects [Palmer, 1975, Biederman et al., 1982,

Original: ,1982; Chun & Jiang, 1998; Henderson & Hollingworth, 1999; Torralba, 2003).
Revised: Chun and Jiang, 1998, Henderson and Hollingworth, 1999, Torralba, 2003].

Original: In addition, spuriousfeatures may appear on the object itself (e.g., facial attributes).
Revised: In addition, spurious features may appear on the object itself (e.g., facial attributes).

Original: Thus discarding the context andbackground may be a futile strategy in these cases (Khani & Liang, 2021).
Revised: Thus discarding the context and background may be a futile strategy in these cases.

Original: In mission-critical applications, relying on a small number of input features may increase the risk offailure.

Original: Recent approaches either use uncertainty estimation (Geifman & El-Yaniv, 2017) or equipmodels with a “rejection” option (Geifman & El-Yaniv, 2019) to reduce such risks.

Original: Instead of requiring contextual and background information to be discarded or relying on a limitednumber of features, we propose a single-epoch ﬁnetuning technique called MaskTune that prevents amodel from learning only the “ﬁrst” simplest mapping (potentially spurious correlations) from theinput to the corresponding target variable.
Revised: Instead of requiring contextual and background information to be discarded or relying on a limited number of features, we propose a single-epoch finetuning technique called MaskTune that prevents a model from learning only the “first” simplest mapping (potentially spurious correlations) from the input to the corresponding target variable.

Original: Our method forces the modeltraining to explore otherinput variables by concealing (masking) the ones that have already been deemed discriminatory.
Revised: MaskTune forces the model to explore other input variables by concealing (masking) the ones that have already been deemed discriminatory.

Original: Aswe ﬁnetune the model with a new set of masked samples, we cause the training to escape its myopicand greedycapacity-preserving feature-seeking approach and encourage exploring and leveragingmore input variables.
Revised: As we finetune the model with a new set of masked samples, we force the training to escape its myopic and greedy feature-seeking approach and encourage exploring and leveraging more input variables.

Original: In other words, as the previous clues are hidden, the model is constrained toﬁnd alternative loss-minimizing input-target mappings.
Revised: In other words, as the previous clues are hidden, the model is constrained to find alternative loss-minimizing input-target mappings.

Original: MaskTune conceals the ﬁrst clues discoveredby a fully trained model, whether they are spurious or not.
Revised: MaskTune conceals the first clues discovered by a fully trained model, whether they are spurious or not.

Original: This forces the model to investigate andleverage new complementary discriminatory input features.
Revised: This forces the model to investigate and leverage new complementary discriminatory input features.

Original: A model relying upon a broader array ofcomplementary features (some may be spurious while others are not) is expected to be more robust totest data missing a subset of these features.
Revised: A model relying upon a broader array of complementary features (some may be spurious while others are not) is expected to be more robust to test data missing a subset of these features.

Original: We then show how our approach can be used for selectiveclassiﬁcation, which is a way of providing models with a rejection option.

Original: Figure 1 visualizes how MaskTune works via a schematic of the cow-on-grass scenario.
Revised: Figure 1 visualizes how MaskTune works via a schematic of the cow-on-grass scenario.

Original: Even inthe absence of spurious correlations, models tend to focus on the shortest path (e.g., ears or skintexture of a cow), which can prevent models from generalizing to scenarios where those speciﬁc partsare missing.
Revised: Even in the absence of spurious correlations, models tend to focus on the shortcut (e.g., ears or skin texture of a cow), which can prevent models from generalizing to scenarios where those specific parts are missing.

Original: However, the object is still recognizable from the remaining parts.
Revised: However, the object is still recognizable from the remaining parts.

Original: As an alternative, MaskTune generates a diverse set of partially masked training examples, forcing the model toinvestigate a wider area of the input features landscape, e.g., new pixels.
Revised: As an alternative, MaskTune generates a diverse set of partially masked training examples, forcing the model to investigate a wider area of the input features landscape e.g., new pixels.

Revised: A further disadvantage of relying on a limited number of features is the model’s inability to know when it does not know.

Revised: Let’s go back to the cow-camel classification example; if cows only appear on grass in the training set then it is unclear which of the grass or the cow refers to the "cow" label.

Revised: A model that only relies on the grass feature can confidently make a wrong prediction when some other object appear in the grass in the test time.

Revised: We need the model to predict only if both cow and grass appear in the picture and abstain otherwise.

Original: To the best of our knowledge, this  is the ﬁrst work to present a ﬁnetuning technique using maskeddata to overcome spurious correlations.
Revised: One method used in the literature to address this issue is selective classification [Geifman and El-Yaniv, 2019, 2017, Khani et al., 2016], which allows a network to reject a sample if it is not confident in its prediction. Selective classification is essential in mission-critical applications such as autonomous driving, medical diagnostics, and robotics as they need to defer the prediction to human if they are uncertain about the prediction. Learning different sets of discriminatory features, in addition to reducing the effect of spurious features, enables MaskTune to be applied to the problem of selective classification. We apply MaskTune to two main tasks: a) robustness to spurious correlations, and b) selective classification.

Revised: We cover four different datasets under (a) including MNIST with synthetic spurious features, CelebA and Waterbirds with spurious features in different subgroups [Sagawa et al., 2019], and the Background Challenge [Xiao et al., 2020] which is a dataset for measuring the reliance of methods on background information for prediction.

Original: {+Under (b) we test MaskTune on CIFAR
+}
Revised: {+Under (b) we test MaskTune on CIFAR
+}

Revised: [Krizhevsky et al., 2009], SVHN [Netzer et al., 2011], and Cats vs. Dogs [Geifman and El-Yaniv, 2019] datasets.

Revised: On both tasks, we outperform or perform similarly to the previous complex methods using our simple technique.

Revised: To the best of our knowledge, this is the first work to present a finetuning technique using masked data to overcome spurious correlations.

Original: Our contributions are summarized as follows:
Revised: Our contributions are summarized as follows:

Original: We propose MaskTune , a new technique to reduce the effect of spurious correlations or over-reliance on a limited number of input features without any supervision such as object location or data subgroup labels.
Revised: We propose MaskTune , a new technique to reduce the effect of spurious correlations or over-reliance on a limited number of input features without any supervision such as object location or data subgroup labels.

Original: We show that MaskTune leads to learning a model that does not rely solely on the initially discovered features.3.
Revised: We show that MaskTune leads to learning a model that does not rely solely on the initially discovered features.

Original: We show how our method can be applied to selective classiﬁcation tasks.
Revised: We show how our method can be applied to selective classification tasks.

Original: We empirically verify the robustness of the learned representations to spurious correlations on a variety of datasets.
Revised: We empirically verify the robustness of the learned representations to spurious correlations on a variety of datasets.

Original:  Method
Revised:  Method

Original: Setup.
Revised: Setup.

Original: We consider the supervised learning setting with inputs x ∈ X ⊂ R d and correspondinglabels y ∈ Y = { 1 , . . .
Revised: We consider the supervised learning setting with inputs x ∈ X ⊂ R d and corresponding labels y ∈ Y = { 1 , . . .

Original: , k } .
Revised: , k } .

Original: We assume having access to samples D 

Revised: We assume having access to samples D

Original: = { ( x i , y i ) } ni =1 drawn from anunknown underlying distribution p data ( x, y ) .
Revised: = { ( x i , y i ) } ni =1 drawn from an unknown underlying distribution p data ( x, y ) .

Original: Our goal is to learn the parameters θ ∈ Θ of a prediction model m θ : X → Y that obtains lowclassiﬁcation error w.r.t some loss function (e.g., cross entropy) (cid:96) : 
Revised: Our goal is to learn the parameters θ ∈ Θ of a prediction model m θ : X → Y that obtains low classification error w.r.t some loss function (e.g., cross entropy) ℓ : Θ × ( X × Y )

Original: Θ × ( X × Y ) → R .
Revised: → R .

Original: Speciﬁcally,we minimize:
Revised: Specifically, we minimize:

Original: [Equation]
Revised: [Equation]

Original: where n is the number of  training data.
Revised: where n is the number of pairs in the training data.

Original: Besides having good prediction accuracy, we aim to develop a model which does not solely rely onspurious or a limited number of input features.
Revised: Besides having good prediction accuracy, we aim to develop a model which does not solely rely on spurious or a limited number of input features.

Original: We propose to mask the input training data to create anew set.
Revised: We propose to mask the input training data to create a new set.

Original: We then ﬁnetune a fully trained  model with the new masked data to reduce over-reliance onspurious or a limited number of features.
Revised: We then finetune ( for only one epoch ) a fully trained ERM model with the new masked data to reduce over-reliance on spurious or a limited number of features.

Revised: The single epoch fine-tuning is done using a small learning rate e.g., the last decayed learning rate that the ERM has used in the first step.

Revised: We found that large learning rates or more than one epoch fine-tuning leads to forgetting the discriminative features learned by the ERM.

Original: Input Masking.
Revised: Input Masking.

Original: A key ingredient of our approach is a masking function G that is applied ofﬂine(i.e. , after full training).
Revised: A key ingredient of our approach is a masking function G that is applied offline (i.e., after full training).

Original: The goal here is to construct a new masked dataset by concealing the mostdiscriminative features in the input discovered by a model after full training.
Revised: The goal here is to construct a new masked dataset by concealing the most discriminative features in the input discovered by a model after full training.

Original: This should encouragethe model to investigate new features with the masked training set during ﬁnetuning.
Revised: This should encourage the model to investigate new features with the masked training set during finetuning.

Original: As for G , weadopt the xGradCAM (Selvaraju et al., 2017), which was originally designed for a visual explanationof deep models by creating rough localization maps based on the gradient of the model loss w.r.t.the output of a desired model layer.
Revised: As for G , we adopt the xGradCAM [Selvaraju et al., 2017], which was originally designed for a visual explanation of deep models by creating rough localization maps based on the gradient of the model loss w.r.t. the output of a desired model layer.

Original: Given an input image of size H × W × C , xGradCAM outputsa localization map A of size H × W × 1 , which shows the contribution of each pixel of the inputimage in predicting the most probable class, i.e., it calculates the loss by choosing the class withhighest logit value (not the true label) as the target class.
Revised: Given an input image of size H × W × C , xGradCAM outputs a localization map A of size H × W × 1 , which shows the contribution of each pixel of the input image in predicting the most probable class, i.e., it calculates the loss by choosing the class with highest logit value (not the true label) as the target class.

Original: After acquiring the localization map, foreach sample ( x i , y i ) , where x i ∈ X and y i ∈ Y , we mask the locations with the most contributionas:
Revised: After acquiring the localization map, for each sample ( x i , y i ) , where x i ∈ X and y i ∈ Y , we mask the locations with the most contribution as:

Original: [Equation]
Revised: ˆ x i =

Original:  where T refers to a thresholding function by the threshold factor τ (i.e., T = 1 A xi ≤ τ ), and (cid:12) denoteselement-wise multiplication.
Revised: T ( A x i ; τ ) ⊙ x i ; A x i = G ( m θ ( x i ) , y i ) (2) where T refers to a thresholding function by the threshold factor τ (i.e., T = 1 A xi ≤ τ ), and ⊙ denotes element-wise multiplication.

Original: As the resolution of A is typically coarser than that of the input data,
Revised: As the resolution of A is typically coarser than that of the input data,

Original: T ( A x i ) is up-sampled to the size of the input.
Revised: T ( A x i ) is up-sampled to the size of the input.

Original: Procedurally, we ﬁrst learn model m initial θ using original unmasked training data D initial .
Revised: Procedurally, we first learn model m initial θ using original unmasked training data D initial .

Original: Then we usem initial θ , G and T to create the masked set D masked .
Revised: Then we use m initial θ , G and T to create the masked set D masked .

Original: Finally, the fully trained predictor m initial θ is tunedusing D masked to obtain m ﬁnal θ .
Revised: Finally, the fully trained predictor m initial θ is tuned using D masked to obtain m final θ .

Original: As for the masking step, any explainability approach can be applied (note that some may have morecomputational complexity, such as ScoreCAM (Wang et al., 2020)).
Revised: As for the masking step, any explainability approach can be applied (note that some may have more computational complexity, such as ScoreCAM [Wang et al., 2020]).

Original: We use xGradCAM (Selvarajuet al., 2017) as it is fast and produces relatively denser heat-maps than other methods (Srinivas &
Revised: We use xGradCAM [Selvaraju et al., 2017] as it is fast and produces relatively denser heat-maps than other methods [Srinivas and

Original: Fleuret, 2019; Selvaraju et al., 2017; Wang et al., 2020).
Revised: Fleuret, 2019, Selvaraju et al., 2017, Wang et al., 2020].

Original: 2.1 MaskTune in Over-parameterized Regimes
Revised: 2.1 MaskTune in Over-parameterized Regimes

Original: Consider the overparametrized regime, in which the model family has sufﬁcient complexity to fully ﬁtthe training data.
Revised: Consider the overparametrized regime, in which the model family has sufficient complexity to fully fit the training data.

Original: It has been shown that deep neural nets are overparametrized and can ﬁt completelyrandom data (Zhang et al., 2021).
Revised: It has been shown that deep neural nets are overparametrized and can fit completely random data [Zhang et al., 2021].

Original: The generalization ability of deep neural nets are still not clear, butthere are some speculation that connect the deep network generalization to their tendency of choosingsimple functions that ﬁt the training data (Valle-Perez et al., 2018; Arpit et al., 2017).
Revised: The generalization ability of deep neural nets is still not clear, but there are some speculation that connect the deep network generalization to their tendency of choosing simple functions that fit the training data [Valle-Perez et al., 2018, Arpit et al., 2017].

Original: However, thissimplicity bias can causes side effects such as their poor performance with respect to adversarialexamples (Raghunathan et al., 2019) or to distribution shifts (Khani & Liang, 2021; Shah et al., 2020).
Revised: However, this simplicity bias can cause side effects such as their poor performance with respect to adversarial examples [Raghunathan et al., 2019] or to distribution shifts [Khani and Liang, 2021, Shah et al., 2020].

Original: Here we study the effect of masking input features on complexity of a model in  situation whereindeed the training procedure chooses the lowest complex model that ﬁt training data.
Revised: Here we study the effect of masking input features on complexity of a model in a situation where indeed the training procedure chooses the least complex model that fits training data.

Original: We show that in this case masking will result in learning a more complex model that discovers new features as theprevious ones are blocked.
Revised: We show that in this case masking will result in learning a more complex model that discovers new features as the previous ones are blocked.

Original: Formally, let C denote a function that measures model complexity and assume that the maskingfunction T (as described in 2) only returns binary values, i.e., an indicator function that only keepssome features and zeros out the rest.
Revised: Formally, let C denote a function that measures model complexity and assume that the masking function T (as described in 2) only returns binary values, i.e., an indicator function that only keeps some features and zeros out the rest.

Original: We show that if training procedure returns the lowest complexmodel then masking results in a more complex model.
Revised: We show that if training procedure returns the least complex model then masking results in a more complex model.

Original: Proposition 1.
Revised: Proposition 1.

Original: Consider an optimizing procedure that ﬁnds min C ( m θ ) , s.t., (cid:96) ( m θ ) = 0 as deﬁnedin 1.
Revised: Consider an optimizing procedure that finds min C ( m θ ) , s.t., ℓ ( m θ ) = 0 as defined in 1.

Original: Let masking function T return binary values.
Revised: Let masking function T return binary values.

Original: If both models m θ and m initial θ ﬁt training data(i.e., zero loss) then we have C ( m ﬁnal θ ) ≥ C ( m initial θ ) .
Revised: If both models m θ and m initial θ fit the training data (i.e., zero loss) then we have C ( m final θ ) ≥ C ( m initial θ ) .

Original: Proof.
Revised: Proof.

Original: Note that both models belong to the model family ( m initial θ , m ﬁnal θ ∈ Θ ), and they both ﬁt thetraining data.
Revised: Note that both models belong to the model family ( m initial θ , m final θ ∈ Θ ), and they both fit the training data.

Original: In the ﬁrst step, training procedure chooses m initial θ over m ﬁnal θ ; therefore according toour assumption C ( m ﬁnal θ ) ≥ C ( m initial θ ) .
Revised: In the first step, training procedure chooses m initial θ over m final θ ; therefore according to our assumption C ( m final θ ) ≥ C ( m initial θ ) .

Original: 2.2 Adapting MaskTune for Selective Classiﬁcation
Revised: 2.2 Adapting MaskTune for Selective Classification

Original: A further disadvantage of relying on a limited number of features is the model’s inability to knowwhen it does not know.

Original: Let’s go back to the cow-camel classiﬁcation example; if cows only appearon the grass in the training set then it is unclear which of the grass or the cow refers to the "cow"label.

Original: A model that only relies on the grass feature can conﬁdently make a wrong prediction whensome other object appear in the grass in the test time.

Original: We need the model to predict only if both cowand grass appear in the picture and abstain otherwise.

Original: One method used in the literature to address this issue is selective classiﬁcation (Geifman & ElYaniv, 2019, 2017), which allows a network to reject a sample if it is not conﬁdent in its prediction.

Original: Selective classiﬁcation is essential in mission-critical applications such as autonomous driving,medical diagnostics, and robotics as they need to defer the prediction to human if they are uncertainabout the prediction.

Original: We now demonstrate how to adapt MaskTuneto be used for the selectiveclassiﬁcation problem.
Revised: Here we show how to use MaskTune for the selective classification problem.

Original: In order to make a more reliable prediction, we ensemble the original model ( m initial θ ) and MaskTune( m ﬁnal θ ) and only predict if both models agree.
Revised: In order to make a more reliable prediction, we ensemble the original model ( m initial θ ) and MaskTune ( m final θ ) and only predict if both models agree.

Original: As a result, if there exist two sets of features that canpredict the label, our method  predicts if both of them predict thesame label (e.g., grass and cow in Figure 1).
Revised: As a result, if there exist two sets of features that can predict the label, our method only predicts if both agree on the label (e.g., grass and cow in Figure 1).

Original: To get an intuition on the performance of MaskTune for selective classiﬁcation, similar to Khaniet al. we analyze the noiseless overparametrized linear regression.
Revised: To get an intuition on the performance of MaskTune for selective classification, similar to [Khani et al., 2016] we analyze the noiseless overparametrized linear regression.

Original: We show that MaskTuneadaptation, explained above, abstains in the presence of covariate shift, thus leading to a more reliableprediction.
Revised: We show that MaskTune adaptation, explained above, abstains in the presence of covariate shift, thus leading to a more reliable prediction.

Original: In particular, we show that MaskTune only predicts if the relationship between masked and unmaskedfeatures in training data holds in the test time.
Revised: In particular, we show that MaskTune only predicts if the relationship between masked and unmasked features in training data holds in the test time.

Original: For example, if features describing “cow” arepredictable of “grass” in training data (i.e., they always co-occur) then we only predict if they cooccur in the test data as well.
Revised: For example, if features describing “cow” are predictable of “grass” in training data (i.e., they always co-occur) then we only predict if they cooccur in the test data as well.

Original: Formally, let s denote the masked feature after the ﬁrst round and zdenote the rest of the d − 1 features.
Revised: Formally, let s denote the masked feature after the first round and z denote the rest of the d − 1 features.

Original: As we are in the overparametrized regime s can be predictedfrom z , let β be the min-norm solution for predicting s from z , i.e., s = βz for training data.
Revised: As we are in the overparametrized regime, s can be predicted from z , let β be the min-norm solution for predicting s from z , i.e., s = β ⊤ z for training data.

Original: We nowshow that both models agree with each other for a new test set if s = βz in the test time as well.
Revised: We now show that both models agree with each other for a new test set if s = β ⊤ z in the test time as well.

Original: Inother words, there is not covariate shift between s and z .
Revised: In other words, there is not covariate shift between s and z .

Original: Proposition 2.
Revised: Proposition 2.

Original: Let S ∈ R n be the concatenation, across all n training samples, of a single maskedscalar feature, out of d possible features, and Z ∈ R n × ( d − 1) be the remaining d − 1 features.
Revised: Let S ∈ R n be the concatenation, across all n training samples, of a single masked scalar feature, out of d possible features, and Z ∈ R n × ( d − 1) be the remaining d − 1 features.

Original: Letthe model family Θ be the linear functions, and the optimization function chooses the min L 2 -normsolution that ﬁts the training data.
Revised: Let the model family Θ be the linear functions, and the optimization function chooses the min L 2 -norm solution that fits the training data.

Original: The models trained with and without S agree on the predictedoutput for a new test set ( z, s ) , iff s =
Revised: The models trained with and without S agree on the predicted output for a new test set ( z, s ) , iff s =

Original: ( Z (cid:62) ( ZZ T ) − 1 S ) (cid:62) z .
Revised: ( Z ⊤ ( ZZ T ) − 1 S ) ⊤ z .

Original: The full proof is in the appendix.
Revised: The full proof is in the appendix.

Original: Note that ( Z (cid:62) ( ZZ T ) − 1 S ) is the min-norm solution for predictings from z in the training data.
Revised: Note that ( Z ⊤ ( ZZ T ) − 1 S ) is the min-norm solution for predicting s from z in the training data.

Original: This proposition states that the models agree only if the relation betweenthe masked feature ( s ) and remaining features ( z ) in training holds in the test data as well.
Revised: This proposition states that the models agree only if the relation between the masked feature ( s ) and remaining features ( z ) in training holds in the test data as well.

Original: In practice, we need to trade off between the coverage and precision.
Revised: In practice, we need to trade off between the coverage and precision.

Original: Therefore, instead of a hardthreshold and predicting  if both models agree, we ensemble m initial θ and m ﬁnal θ by multiplying their
Revised: Therefore, instead of a hard threshold and predicting only if both models agree, we ensemble m [Table] θ [Table]

Original: Input
Revised: by [Table]

Original: ERM probabilities
Revised: their probabilities.

Original: and we ﬁnd the desired threshold for abstaining in the validation setin order to achievethe coverage goal (see Section 3 for details).  
Revised: In order to achieve the coverage goal, we find the desired threshold for abstaining in the validation set (see Section 4 for details).  Implementation Details

Revised: Classification with Spurious Features.

Original: Experimental Results
Revised: For all of the datasets in Section 4.1, we use SGD with a momentum of 0.9 and weight decay of 1e-4. For biased MNIST, we used a simple convolutional neural network with four convolutional layers and a linear head. We trained this model with a batch size of 128 for 100 epochs with a learning rate of 0.01. We decreased the learning rate by a factor of 0.5 every 25 epochs.

Revised: In the case of Background Challenge, we trained an ImageNet pre-trained

Revised: ResNet50 with a batch size of 1,024 for 100 epochs decaying the learning rate by 0.1 after everyepochs.

Revised: For the CelebA dataset, we trained an ImageNet pre-trained ResNet50 with a batch size of 512 for 20 epochs with a learning rate of 1e-4.

Revised: Selective Classification.

Revised: For all the datasets in Section 4.2, we trained a ResNet32 from scratch.

Revised: For

Revised: CIFAR-10, we used SGD with a momentum of 0.9, weight decay of 1e-4, learning rate of 0.1, and batch size of 128.

Revised: We trained it for 300 epochs and halved the learning rate every 25 epochs.

Revised: For

Revised: SVHN, we used the same hyperparameters as CIFAR-10 with only two differences: the weight decay and batch size were 5e-4 and 1,024, respectively.

Revised: For Cats vs. Dogs dataset, we used Adam optimizer with a weight decay of 1e-4 and a learning rate of 0.001.

Revised: We trained the model for 200 epochs with a batch size of 128 and dropped the learning rate by a factor of 0.1 on epoch 50.

Revised: Masking Threshold.

Revised: Our goal is to mask the most important features, i.e., the core of the heat maps generated by the explainability methods.

Revised: Masking a few input variables has almost no effect on the model’s behavior, whereas large masks may destroy useful signals in the input, resulting in very low training accuracy.

Revised: To reduce the search space of the masking threshold τ across all tasks, we experimented with τ = { µ i , µ i + σ i , µ i + 2 σ i , µ i + 3 σ i } where µ i and σ i represent mean and standard deviation over the heatmap values for training sample x i .

Revised: We also experimented with mean value over all training samples, using soft masks (i.e., no threshold), and sorting and masking the

Original:  We evaluate our MaskTune method on two main applications: a) classiﬁcation with spuriouscorrelations—we expect MaskTune to prevent such correlations by identifying and masking them andb) selective classiﬁcation where the ability to abstain is critical—we expect MaskTune to improvereliability by forcing a model to investigate additional variables in the input and learn more complexrelationships.
Revised: K-top activated variables. We found µ i + 2 σ i works better in general as µ i and µ i + σ i remove a large portion of the input while µ i + 3 σ i removes very few variables.  Experimental Results We evaluate our MaskTune method on two main applications: a) classification with spurious correlations—we expect MaskTune to prevent such correlations by identifying and masking them and b) selective classification where the ability to abstain is critical—we expect MaskTune to improve reliability by forcing a model to investigate additional variables in the input and learn more complex relationships.

Original: In each experiment, we compare our method to relevant baselines and competingmethods.
Revised: In each experiment, we compare our method to relevant baselines and competing methods. [Figure]

Original: 3.1 Classiﬁcation with Spurious Features
Revised: 4.1 Classification with Spurious Features

Original: MNIST with Colored Squares on the Corners: As a warm-up, we test the ability of our methodto distinguish between two MNIST digit groups (0-4 and 5-9) in the presence of spurious features.
Revised: MNIST with Colored Squares on the Corners: As a warm-up, we test the ability of our method to distinguish between two MNIST digit groups (0-4 and 5-9) in the presence of spurious features.

Original: We construct a dataset such that a classiﬁer would achieve poor accuracy by relying on spuriousinput features.
Revised: We construct a dataset such that a classifier would achieve poor accuracy by relying on spurious input features.

Original: We place 99% and 1% ofall digits labelled "0-4" and "5-9" in the training set on abackground with a small blue square on the top left corner, respectively, and keep the remaining dataintact.
Revised: To this end, we group digits labeled 0 to 4 into class 0 and those labeled 5 to 9 into class 1. Next, in the training set, we place 99% and 1% of the new class 0 and new class 1 data on a background with a small blue square on the top left corner, respectively, and keep the remaining data intact.

Original: We use two test sets during testing: the original raw MNIST test set and a biassed test set.
Revised: We use two test sets during testing: the original raw MNIST test set and a biased test set.

Original: Tocreate the biased test set we place all of the digits from the group “5-9" on a background with a smallblue square (representing spurious features) on the top left corner and keep all of the digits fromthe group “0-4" unchanged (i.e., we don’t add any squares to their background and use the originalimages).
Revised: To create the biased test set we place all of the digits from the group “5-9” on a background with a small blue square (representing spurious features) on the top left corner and keep all of the digits from the group “0-4” unchanged (i.e., we don’t add any squares to their background and use the original images).

Original: Figure 2 demonstrates the performance of the ERM, RandMask, and our MaskTune on bothoriginal and biased test sets.
Revised: Figure 2 demonstrates the performance of the ERM, RandMask, and our MaskTune on both original and biased test sets.

Original: The RandMask method is similar to MaskTune , but masks a randomlychosen area of the input image.
Revised: The RandMask method is similar to MaskTune , but masks a randomly chosen area of the input image.

Original: As shown in Figure 2, MaskTune outperforms other methods by alarge margin on boththe test sets.
Revised: As shown in Figure 2, MaskTune outperforms other methods by a large margin on both test sets.

Original: As demonstrated in Figure 2 (right), MaskTune forces the model toexplore more features, as opposed to the ERM, which only looks into the spurious feature (the bluesquare).
Revised: As demonstrated in Figure 2 (right), MaskTune forces the model to explore more features, as opposed to the ERM, which only looks into the spurious feature (the blue square). We also ran an experiment with multiple spurious features (Appendix A) and reported results for iterative version of MaskTune .

Original: Classiﬁcation with Spurious Features in Subgroups.
Revised: Classification with Spurious Features in Subgroups.

Original: In this experiment, we leverage the
Revised: In this experiment, we leverage the

Original: CelebA (Liu et al., 2015) and the Waterbirds (Sagawa et al., 2019) datasets.
Revised: CelebA [Liu et al., 2015] and the Waterbirds [Sagawa et al., 2019] datasets.

Original: In the CelebA dataset,there is a high correlation between features gender={male, female} and hair_color={blond, dark},meaning that the feature gender is used as a proxy to predict the hair_color.
Revised: In the CelebA dataset, there is a high correlation between features gender={male, female} and hair_color={blond, dark}, meaning that the feature gender might be used as a proxy to predict the hair_color.

Original: In other words, an ERM would assign the label dark to male images since the majority of male images have dark hair,and there are only 1,387 (0.85%) blond males in the training set of size 162,770.
Revised: In other words, an ERM would assign the label dark to male images since the majority of male images have dark hair, and there are only 1,387 (0.85%) blond males in the training set of size 162,770.

Original: To balance theaccuracy of predictions on different subgroups of data, the existing methods (Sagawa et al., 2019)use privileged subgroup information, i.e., wrong predictions of an ERM on the worst group (blondmales) can be penalized during training.
Revised: To balance the accuracy of predictions on different subgroups of data, the existing methods [Sagawa et al., 2019] use subgroup information, i.e., wrong predictions of an ERM on the worst group (blond males) can be penalized during training.

Original: Some other approaches (Levy et al., 2020; Nam et al., 2020; Pezeshki et al., 2021; Liu et al., 2021; Taghanaki et al., 2021) use such information during model Input ERM
Revised: Some other approaches [Levy et al., 2020, Nam et al., 2020, Pezeshki et al., 2021, Taghanaki et al., 2021, Liu et al., 2021] use subgroup information during model selection (labeled validation set).

Original: MaskTune

Original: Input

Original: ERM

Original: MaskTune selection (labeled validation set).

Original: However, it is impractical to recognize “all” subgroups in a datasetand label them.
Revised: However, it is impractical to recognize “all” subgroups in a dataset and label them.

Original: On the other hand, our MaskTune uses group information neither during trainingnor during model selection. In Table 1, we compare our MaskTune  to both these groups of methods.
Revised: In Table 1, we show that MaskTune achieves comparable performance to both these groups of methods without using group information in training or in model selection.

Original:  MaskTune achieves comparable performance to the methods which use group level supervision whilesigniﬁcantly outperforming (78% vs. 55% classiﬁcation accuracy) the methods that do not use anysupervision on the worst-group accuracy. [Table]
Revised: We also show MaskTune significantly improves worst-group accuracy (78% vs. 55% classification accuracy) in comparison to methods that do not use subgroup information during training or model selection.

Original: No [Table] [Table] [Table]
Revised: In Figure 3 (left), we highlight the important input features for predicting hair color on the CelebA dataste.

Original: No [Table] [Table] No
Revised: As shown, the ERM model leverages gender features while MaskTune forces to investigate other features as well.

Original: [Table] DivDis (Lee et al., 2022) No No 55.
Revised: The Waterbirds dataset [Sagawa et al., 2019] was proposed to assess the degree to which models pick up spurious correlations in the training set.

Original: MaskTune [Table]
Revised: We discovered and fixed two issues with the Waterbirds dataset: a) because the background images in the Places dataset [Zhou et al., 2017]

Original: ours) No
Revised: may already

Original: No 78.0 ± 1. 91.3 ± 0. In Figure 3, we highlight the important input features for predicting hair color. As shown, the ERMmodel leverages gender features while MaskTune forces to investigating other features as well.
Revised: Method

Original: We discovered and ﬁxed two problems with the Waterbirds dataset: a) Because the backgroundimages in the Places dataset (Zhou et al., 2017) may already contain bird images, multiple birdsmay appear in an image after overlaying the segmented bird images from the Caltech-UCSD Birds200-2011 (CUB) dataset
Revised: [Table] Group DRO Sagawa et al. [Table]

Original: (Wah et al., 2011).
Revised: CVaR DRO Levy et al. [Table] LfF Nam et al. [Table]

Revised: SD Pezeshki et al.

Revised: [Table]

Revised: JTT Liu et al.

Revised: [Table]

Revised: CIM Taghanaki et al.

Revised: [Table]

Revised: [Table]

Revised: ERM Sagawa et

Revised: [Table]

Revised: [Table]

Revised: CVaR DRO Levy et al.

Revised: [Table]

Revised: LfF Nam et al.

Revised: [Table]

Revised: JTT Liu et al.

Revised: [Table]

Revised: DivDis Lee et al.

Revised: [Table]

Revised: MaskTune (ours)

Revised: [Table]

Revised: Input

Revised: [Table]

Revised: contain bird images, multiple birds may appear in an image after overlaying the segmented bird images from the Caltech-UCSD Birds-200-2011 (CUB) dataset

Revised: [Wah et al., 2011].

Original: For example, the label of an image may be “landbird”,but the image contains both land and water birds.
Revised: For example, the label of an image may be “landbird”, but the image contains both land and water birds.

Original: Such images were removed  from the dataset.
Revised: We manually removed such images from the dataset.

Original: b) Because the names of the species are similar, some land birds have been mislabeled as waterbirds.
Revised: b) Because the names of the species are similar, some land birds have been mislabeled as waterbirds which we corrected.

Revised: The corrected Waterbirds dataset can be found on MaskTune ’s GitHub page.

Original: After addressing these issues, the ERM model’s worst-group accuracy increased from 60% which isreported in (Sagawa et al., 2019) to 82%.
Revised: After addressing the two issues, the ERM model’s worst-group accuracy increased from 60% which is reported in [Sagawa et al., 2019] to 80.8 ± 1.3%.

Original: We repeated the group-DRO method and ERM experimentson the corrected waterbirds dataset and reported the results in Table 2.
Revised: We repeated the group-DRO method and ERM experiments on the corrected waterbirds dataset and reported the results in Table 2.

Original: As demonstrated, our method(without any group supervision) achieves similar accuracy to the group-DRO that beneﬁts from fullsupervision.
Revised: As demonstrated, our method (without any group supervision) achieves similar accuracy to the group-DRO that benefits from full supervision.

Original: In Figure 3, we visualized feature importance before and after applying our MaskTune .
Revised: In Figure 3 (right), we visualized feature importance before and after applying our MaskTune on the modified Waterbirds dataset.

Original: [Table]
Revised: [Table]

Original: The Background Challenge.
Revised: The Background Challenge.

Original: As a further step, we evaluated MaskTune on the Background Challenge data (Xiao et al., 2020) to see if the positive observations from the MNIST experiment apply toa more realistic scenario.
Revised: As a further step, we evaluated MaskTune on the Background Challenge data [Xiao et al., 2020] to see if the positive observations from the MNIST experiment apply to a more realistic scenario.

Original: The Background Challenge is a publicly available dataset that consists of
Revised: The Background Challenge is a publicly available dataset that consists of

Original: ImageNet-9 (Deng et al., 2009) test sets with various levels of foreground and background signals.
Revised: ImageNet-9 [Deng et al., 2009] test sets with various levels of foreground and background signals.

Original: Itis intended to assess how much deep classiﬁers rely on spurious features for image classiﬁcation.
Revised: It is intended to assess how much deep classifiers rely on spurious features for image classification.

Original: Weused two conﬁgurations to compare MaskTune ’s performance: Only FG, in which the background iscompletely removed, Mixed-same, in which the foreground is placed on  different background fromthe same class, and Mixed-rand, where the foreground is overlaid onto a random background.
Revised: We used two configurations to compare MaskTune ’s performance: Only FG, in which the background is completely removed, Mixed-same, in which the foreground is placed on a different background from the same class, and Mixed-rand, where the foreground is overlaid onto a random background.

Original: As shown in Table 3, MaskTune outperforms the baseline ResNet-50’s performance by 2.5% onthe Only-FG test set and 1.2% on Mixed-same, showing that MaskTune does not rely much on thebackground and uses both background and foreground for prediction. 
Revised: As shown in Table 3, MaskTune outperforms the baseline ResNet-50’s performance by 2.5% on the Only-FG test set and 1.2% on Mixed-same, showing that MaskTune does not rely much on the background and uses both background and foreground for prediction. On the Mixed-same and Only-FG test sets

Revised: , MaskTune outperforms other techniques because mixing/removing the background texture/info confuses other methods.

Original: These results show that ourtechnique helps to learn task-relevant features without depending on nuisance signal sources.
Revised: These results show that our technique helps to learn task-relevant features without depending on nuisance signal sources.

Original: [Table]
Revised: [Table]

Original: For selective classiﬁcation task, we evaluatedour MaskTune on three datasets: CIFAR-10 (Krizhevskyet al., 2009), SVHN (Netzer et al., 2011), and Cats vs. Dogs (Geifman & El-Yaniv, 2019), withcoverage values of {90%, 95%, 100%}.
Revised: For selective classification task, we evaluated MaskTune on three datasets: CIFAR-10 [Krizhevsky et al., 2009], SVHN [Netzer et al., 2011], and Cats vs. Dogs [Geifman and El-Yaniv, 2019], with coverage values of {90%, 95%, 100%}.

Original: For example 90% coverage means abstaining 10% of thesamples.
Revised: For example 90% coverage means abstaining 10% of the samples.

Original: Given input image j and the original ( O ) and ﬁnetuned ( M ) models, let their respective inference-timeprediction probabilities for class i be P Mij and P Oij .
Revised: Given input image j and the original ( m initial θ ) and finetuned ( m final θ ) models, let their respective inference-time prediction probabilities for class i be P init ij and P final ij .

Original: Then MaskTune does not abstain and declares theclass i∗ as the predicted class, if and only if P Mij · P Oij > γ , where γ is a threshold that allows forcontrolling the target coverage.
Revised: Then MaskTune does not abstain and declares the class i as the predicted class, iff P init ij · P final ij > γ , where γ is a threshold that allows for controlling the target coverage.

Original: To ﬁnd the proper γ for achieving the targeted coverage, we iterateover the validation set and ﬁnd a threshold that, if applied to the validation set, would give the desired coverage, then use the found threshold in the test time (i.e., if our target coverage is 90%, we seek fora threshold which if we apply for abstention on the validation set, we will abstain predicting 10% ofvalidation data).
Revised: To find the proper γ for achieving the targeted coverage, we iterate over the validation set and find a threshold that, if applied to the validation set, would give the desired coverage (i.e., if our target coverage is 90%, we seek for a threshold which if we apply for abstention on the validation set, we will abstain predicting 10% of validation data).

Original: Note that the probabilities of the two models are multiplied to ensure that as thevalue of either probability decreases, the possibility of abstention increases.
Revised: Note that the probabilities of the two models are multiplied to ensure that as the value of either probability decreases, the possibility of abstention increases.

Original: We compare our method to the other selective classiﬁcation approaches such as Softmax Response(SR)
Revised: We compare our method to the other selective classification approaches such as Softmax Response (SR)

Original: (Geifman & El-Yaniv, 2019), SelectiveNet (SN) (Geifman & El-Yaniv, 2019), Deep Gamblers(DG)
Revised: [Geifman and El-Yaniv, 2019], SelectiveNet (SN) [Geifman and El-Yaniv, 2019], Deep Gamblers (DG)

Original: (Liu et al., 2019), and One-sided Prediction (OSP) (Gangrade et al., 2021).
Revised: [Liu et al., 2019], and One-sided Prediction (OSP) [Gangrade et al., 2021].

Original: Like the OSP, weﬁrst split the training data into train and validation sets.
Revised: Like the OSP, we first split the training data into train and validation sets.

Original: After training, we search for the abstentionthreshold on the validation set, then use it at the test time.
Revised: After training, we search for the abstention threshold on the validation set, then use it at the test time.

Original: As seen in Table 4, MaskTune outperformsall prior techniques on all three datasets on 95% and 100% coverage rates.
Revised: As seen in Table 4, MaskTune outperforms all prior techniques on all three datasets on 95% and 100% coverage rates.

Original: For results on morecoverage rates refer to Appendix A.
Revised: For results on more coverage rates refer to Appendix A.

Original: [Table]
Revised: [Table]

Original: Implementation Details

Original: Classiﬁcation with Spurious Features.

Original: For all of the datasets in Section 3.1, we use SGD witha momentum of 0.9 and weight decay of 1e-4.

Original: For biased MNIST, we used a simple convolutionalneural network with four convolutional layers and a linear head.

Original: We trained this model with a batchsize of 128 for 100 epochs with a learning rate of 0.01.

Original: We decreased the learning rate by a factorof 0.5 every 25 epochs.

Original: In the case of Background Challenge, we trained an ImageNet pre-trained

Original: ResNet50 with a batch size of 1024 for 100 epochs decaying the learning rate by 0.1 after every 30epochs.

Original: For the CelebA dataset, we trained an ImageNet pre-trained ResNet50 with a batch size of512 for 20 epochs with a learning rate of 1e-4.

Original: Selective Classiﬁcation.

Original: For all the datasets in Section 3.2, we trained a ResNet32 from scratch.

Original: For CIFAR-10, we used SGD with a momentum of 0.9, weight decay of 1e-4, learning rate of 0.1,and batch size of 128.

Original: We trained it for 300 epochs and halved the learning rate every 25 epochs.

Original: For

Original: SVHN, we used the same hyperparameters as CIFAR-10 with only two differences: the weight decayand batch size were 5e-4 and 1024, respectively.

Original: For Cats vs. Dogs dataset, we used Adam optimizerwith a weight decay of 1e-4 and a learning rate of 0.001.

Original: We trained the model for 200 epochs with abatch size of 128 and dropped the learning rate by a factor of 0.1 on epoch 50.

Original: Masking Threshold.

Original: Our goal is to mask the most important features, i.e., the core of the heatmaps generated by the explainability methods.

Original: Masking a few input variables has almost no effecton the model’s behavior, whereas large masks may destroy useful signals in the input, resulting invery low training accuracy.

Original: To reduce the search space of the masking threshold τ across all tasks,we experimented with τ = { µ i , µ i + σ i , µ i + 2 σ i , µ i + 3 σ i } where µ i and σ i represent mean andstandard deviation over the heatmap values for training sample x i .

Original: We found µ i + 2 σ i works best asµ i and µ i + σ i remove a large portion of the input

Original: while µ i + 3 σ i removes very few variables.

Original: Hardware.

Original: We used two multi-GPU machines for all the experiments: a) 8 NVIDIA V100 Tensor

Original: [-Core GPUs with 32 GB of memory each, 728 GB RAM, b) 4 NVIDIA Quadro RTX 5000 with 
-]
Revised: [-Core GPUs with 32 GB of memory each, 728 GB RAM, b) 4 NVIDIA Quadro RTX 5000 with 
-]

Original: GB memory each, and 128 GB RAM.

Original:  Related Work
Revised:  Related Work

Original: In the following, we review the related works which have attempted to mitigate the effect of spuriouscorrelation in training deep models.
Revised: In the following, we review the related works which have attempted to mitigate the effect of spurious correlation in training deep models.

Original: We also discuss previous methods that use attention-based onlinemasking and how our method differs.
Revised: We also discuss previous methods that use attention-based online masking and how our method differs.

Original: Robustness to Spurious Correlations.
Revised: Robustness to Spurious Correlations.

Original: Distributionally robust optimization (DRO) (Ben-Tal et al. ,2013; Gao et al., 2017; Duchi et al., 2021) has been proposed to improve generalization to worst cases(minority distributions) in a dataset.
Revised: Distributionally robust optimization (DRO) [Ben-Tal et al., 2013, Gao et al., 2017, Duchi et al., 2021] has been proposed to improve generalization to worst cases (minority distributions) in a dataset.

Original: However, the DRO objective leads to disproportionate attentionto the worst cases, even if they are implausible.
Revised: However, the DRO objective leads to disproportionate attention to the worst cases, even if they are implausible.

Original: To address this issue, Sagawa et al.
Revised: To address this issue, Sagawa et al.

Original: (2019) proposedto leverage subgroup information during optimization.
Revised: [2019] proposed to leverage subgroup information during optimization.

Original: Although this method reduces the likelihoodof the worst-case failure, it is based on prior solid information (i.e., subgroup labels), which is notalways available.
Revised: Although this method reduces the likelihood of the worst-case failure, it is based on prior solid information (i.e., subgroup labels), which is not always available.

Original: Several efforts have been made to reduce the subgroup-level supervision.
Revised: Several efforts have been made to reduce the subgroup-level supervision.

Original: Sohoniet al.
Revised: Sohoni et al.

Original: (2020) proposed a clustering-based method for obtaining sub-group information to be used in
Revised: [2020] proposed a clustering-based method for obtaining sub-group information to be used in

Original: DRO setup.
Revised: DRO setup.

Original: However, determining the number of clusters (sub-groups) in a dataset is not trivial sincea small number may still dismiss minor subgroups while a large number may lower overall accuracy.
Revised: However, determining the number of clusters (sub-groups) in a dataset is not trivial since a small number may still dismiss minor subgroups while a large number may lower overall accuracy.

Original: Yaghoobzadeh et al.
Revised: Yaghoobzadeh et al.

Original: (2019) proposed measuring sample accuracy throughout training to discoverforgettable instances, then ﬁne-tuning models using such samples to increase model resilience againstspurious correlations.
Revised: [2019] proposed measuring sample accuracy throughout training to discover forgettable instances, then fine-tuning models using such samples to increase model resilience against spurious correlations.

Original: Chen et al.
Revised: Chen et al.

Original: (2020) demonstrated that self-training can also help decrease theeffect of spurious features, but only if the source classiﬁer is very accurate and there are not toomany isolated sub-groups in the data.
Revised: [2020] demonstrated that self-training can also help decrease the effect of spurious features, but only if the source classifier is very accurate and there are not too many isolated sub-groups in the data.

Original: Liu et al.
Revised: Liu et al.

Original: (2021) discovered that training a model twice helps itbecome resistant to spurious correlation.
Revised: [2021] discovered that training a model twice helps it become resistant to spurious correlation.

Original: However, if a model reaches high classiﬁcation accuracyin the ﬁrst run, this technique is no longer useful since there is not enough misclassiﬁed data toretrain the model.
Revised: However, if a model reaches high classification accuracy in the first run, this technique is no longer useful since there is not enough misclassified data to retrain the model.

Original: MaskTune , on the other hand, employs an orthogonal approach to learn robustrepresentations through gradient-based masking and ﬁne-tuning and does not require subgroup-levellabels during training or model selection.
Revised: MaskTune , on the other hand, employs an orthogonal approach to learn robust representations through gradient-based masking and fine-tuning and does not require subgroup-level labels during training or model selection.

Original: Online Attention-based Masking.
Revised: Online Attention-based Masking.

Original: A large number of studies suggest using attention-based (soft)masks to eliminate irrelevant information from input data during training (Sharma et al., 2015; Wanget al., 2017; Xu et al., 2015; Zheng et al., 2017).
Revised: A large number of studies suggest using attention-based (soft) masks to eliminate irrelevant information from input data during training [Sharma et al., 2015, Wang et al., 2017, Xu et al., 2015, Zheng et al., 2017].

Original: Although these methods increase overall classiﬁcationaccuracy, they are incapable of disregarding spurious correlations since attention may be on an areaof the input where the spurious correlations are dominant during training.
Revised: Although these methods increase overall classification accuracy, they are incapable of disregarding spurious correlations since attention may be on an area of the input where the spurious correlations are dominant during training.

Original: For concentrating attentiononly on the foreground, Li et al.
Revised: For concentrating attention only on the foreground, Li et al.

Original: (2018) developed a guided attention model.
Revised: [2018] developed a guided attention model.

Original: Their technique,however, needs additional annotations of object locations/masks.
Revised: Their technique, however, needs additional annotations of object locations/masks.

Original: Nonetheless, none of these methodsaimed to reduce the effects of spurious correlations.
Revised: Nonetheless, none of these methods aimed to reduce the effects of spurious correlations.

Original: Even with complete supervision (i.e., maskingthe whole background), spurious correlations might still occur in the foreground.
Revised: Even with complete supervision (i.e., masking the whole background), spurious correlations might still occur in the foreground.

Original: MaskTune , on theother hand, produces masks from a fully trained model and uses them for a single epoch ﬁne-tuning.
Revised: MaskTune , on the other hand, produces masks from a fully trained model and uses them for a single epoch fine-tuning.

Original: MaskTune does not require any additional supervision, such as object location annotations.
Revised: MaskTune does not require any additional supervision such as object location annotations.

Original:  Conclusion
Revised:  Conclusion

Original: In this work, we considered the problem of preventing models from learning spurious correlations.
Revised: In this work, we considered the problem of preventing models from learning spurious correlations.

Original: We introduced a new ﬁne-tuning technique that is designed explicitly for spurious correlations.
Revised: We introduced a new fine-tuning technique that is designed explicitly for spurious correlations.

Original: Itenforces a model to explore more variables in the input and map them to the same target.
Revised: It enforces a model to explore more variables in the input and map them to the same target.

Original: Throughexperiments and theoretical analysis on classiﬁcation with nuisance background information whichtypically suffers from the presence of spurious correlations in the data, we showed that models trainedwith MaskTune outperform previous relevant methods.
Revised: Through experiments and theoretical analysis on classification with nuisance background information which typically suffers from the presence of spurious correlations in the data, we showed that models trained with MaskTune outperform previous relevant methods.

Original: We also showed that MaskTune helps toimprove the accuracy signiﬁcantly in selective classiﬁcation tasks.
Revised: We also showed that MaskTune helps to improve the accuracy significantly in selective classification tasks.

Original: Adapting MaskTune for non-imagedata, such as sentiment analysis, can be an intriguing future work.
Revised: Adapting MaskTune for non-image data, such as sentiment analysis, can be an intriguing future work.

Original: Another area for future work ishow best to identify and use the most effective masked samples (rather than all) when ﬁne-tuning theﬁnal model, perhaps using uncertainty information or active learning.
Revised: Another area for future work is how best to identify and use the most effective masked samples (rather than all) when fine-tuning the final model, perhaps using uncertainty information or active learning.

Original: Appendix
Revised: Although MaskTune can help reduce the effect of many types of spurious correlation, such as texture, color, and localized nuisance features e g., artifacts added to x-ray images by medical imaging devices, there are some cases where MaskTune may not be effective, such as a small transformation in all pixel values for some of the images in a dataset. This occurs in medical devices or cameras that add almost imperceptible color (values) to captured images.

Original:  A.1 More experimental results
Revised: A Appendix A.1 More Experimental Results

Original: Selective classiﬁcation.
Revised: Selective Classification.

Original: In table 5 the results of the selective classiﬁcation is depicted for the
Revised: In table 7 we show the results of the selective classification for the

Original: MaskTune for different coverages.
Revised: MaskTune on different coverage values.

Original: We have run every experiment three times and reported the meanand std for each.
Revised: We have run every experiment three times and reported the mean and std for each.

Revised: Multiple Spurious Features Scenario.

Revised: Running MaskTune for only one iteration on a dataset with more than one spurious feature still performs better than ERM.

Revised: If we run MaskTune iteratively, the performance improves even more.

Original: B Missing proofs
Revised: We ran two iterative versions of MaskTune : accumulative and non-accumulative on the MNIST dataset (Table 5). We added two coloured patches to MNIST digits as two distinct spurious features. We ran 1, 2, and 3 iterations of masking. MaskTune works well when we do iterative accumulative masking (i.e., add new masks to the previously masked samples), but when we apply masks from each iteration to the raw input (non-accumulative), the results are not as good as the accumulative version.

Revised: To reduce the running time, we only used one iteration of masking.

Revised: One way to stop the accumulative masking is to monitor the training accuracy.

Revised: If the model is not able to fit the data after a certain number of masking iterations (because there are no useful features left), the training can be stopped.

Revised: [Table]

Revised: 76 ± 0.

Revised: Is it harmful to use MaskTune when we are unsure whether there is a spurious correlation in our data (e.g., ImageNet)?

Revised: Even if there are no spurious features in the input or detecting the spurious feature is difficult,

Revised: MaskTune masks the initially found most discriminative features, resulting in the discovery of the second discriminative set of features.

Revised: For example, if an ERM model initially finds the cow’s head to be the most discriminative feature, masking it (which is not spurious) forces the model to find the second best feature as well, which could be the texture of the cow’s skin.

Revised: The final model would learn both the cow’s head and skin texture as discriminative features, rather than just the head.

Revised: We ran MaskTune on ImageNet.

Revised: The overall test accuracy did not decrease significantly i.e., from 78.

Revised: (ERM) to 78.213 ( MaskTune ).

Revised: After applying MaskTune , classification accuracy on certain classes improved significantly e.g., by over 50%.

Revised: For example for class “projectile, missile” the accuracy was improved from 25% to 80%, for class “crane” from 44% to 82%, for class “skunk, polecat, wood pussy” from 49% to 69%, for class “coffeepot” from 48% to 64%, etc.

Revised: More Aggressive Random Masking as a Baseline.

Revised: We compare MaskTune to a another version of random masking with different window sizes on the MNIST dataset.

Revised: For each image, we randomly selected a window of { 2 × 2 , 3 × 3 , 4 × 4 , ....n × n } pixels where n is the image size divided by 2 .

Revised: As shown in Table 6, MaskTune still outperforms the random masking approach.

Revised: [Table]

Revised: [Table]

Revised: B Hardware

Revised: We used two multi-GPU machines for all the experiments: a) 8 NVIDIA V100 Tensor Core GPUs with 32 GB of memory each, 728 GB RAM, b) 4 NVIDIA Quadro RTX 5000 with 16 GB memory each, and 128 GB RAM.

Revised: C Missing proofs

Original: Proposition 2.
Revised: Proposition 2.

Original: Let S ∈ R n be the concatenation, across all n training samples, of a single maskedscalar feature, out of d possible features, and Z ∈ R n × ( d − 1) be the remaining d − 1 features.
Revised: Let S ∈ R n be the concatenation, across all n training samples, of a single masked scalar feature, out of d possible features, and Z ∈ R n × ( d − 1) be the remaining d − 1 features.

Original: Letthe model family Θ be the linear functions, and the optimization function chooses the min L 2 -normsolution that ﬁts the training data.
Revised: Let the model family Θ be the linear functions [Table] and the optimization function chooses the min L 2 -norm solution that fits the training data.

Original: The models trained with and without S agree on the predictedoutput for a new test set ( z, s ) , iff s =
Revised: The models trained with and without S agree on the predicted output for a new test set ( z, s ) , iff s =

Original: ( Z (cid:62) ( ZZ T ) − 1 S ) (cid:62) z .
Revised: ( Z ⊤ ( ZZ T ) − 1 S ) ⊤ z .

Original: Proof.
Revised: Proof.

Original: Let Y ∈ R
Revised: Let Y ∈ R

Original: d be the targets in training data.
Revised: d be the targets in training data.

Original: Let ˆ θ + s denote the recovered coefﬁcient for theﬁrst d − 1 features and ˆ w be the coefﬁcient for the feature s .[Table]
Revised: Let ˆ θ + s denote the recovered coefficient for the first d − 1 features and ˆ w be the coefficient for the feature s .

Original: The parameters after the ﬁrst round can be obtained by the following optimization problem.
Revised: The parameters after the first round can be obtained by the following optimization problem.

Original: [Equation]
Revised: [Equation]

Original: where ˆ w denotes the weight associate with the feature s .
Revised: where ˆ w denotes the weight associated with the feature s .

Original: Recall that we assumed that after training the ﬁrst model we mask the feature s .
Revised: Recall that we assumed that after training the first model we mask the feature s .

Original: The parameters forthe second model that does not use feature s can be obtained by solving the following optimizationproblem, where ˆ θ − s denote the recovered coefﬁcients. [Equation]
Revised: The parameters for the second model that does not use feature s can be obtained by solving the following optimization problem, where ˆ θ − s denote the recovered coefficients.

Original: [Equation]
Revised: [Equation]

Original: The minimum norm solution in the ﬁrst round is:
Revised: The minimum norm solution in the first round is:

Original: [Equation]
Revised: [Equation]

Original: The min-norm solution for the second round is:
Revised: The min-norm solution for the second round is:

Original: [Equation]
Revised: [Equation]

Original: Now we show that for a new test set ( z (cid:63) , s (cid:63) ) , the models m 1 and m 0 predict the same if there is nocovariate shift between z and s .
Revised: Now we show that for a new test set ( z ⋆ , s ⋆ ) , the models m final θ and m initial θ predict the same if there is no covariate shift between z and s .

Original: We are assuming that ˆ w is not zero meaning that the ﬁrst model isusing the feature s .
Revised: We are assuming that ˆ w is not zero meaning that the first model is using the feature s .

Original: [Equation]
Revised: [Equation]

Revised: For proving the other direction, if we have ( Z ⊤ ( ZZ ⊤ ) − 1 S )) ⊤ z

Revised: ⋆ = s

Revised: ⋆ then we can deduce from

Revised: Equation 11 to Equation 9.