\begin{table}[h!]
\centering
\setlength{\tabcolsep}{3pt}
\caption{Compare \textit{MaskTune} with random masking.}
\begin{tabular}{c|c|c}
Test data type & Masking method & Accuracy     \\ \midrule
biased         & random masking & 59.06$\pm$09.25 \\
original       & random masking & 91.02$\pm$01.87 \\
biased         & MaskTune       & \textbf{98.26$\pm$00.27} \\
original       & MaskTune       & \textbf{98.46$\pm$00.21}
\end{tabular}
\label{tab:randmask2}
\end{table}
\section{Method}
\textbf{Setup.} 
We consider the supervised learning setting with inputs $x \in \mathcal{X} \subset \mathbb{R}^d$ and corresponding labels $y \in \mathcal{Y} = \{1, \ldots, k\}$. We assume having access to samples $\mathcal{D}^0 = \{(x_i, y_i)\}_{i=1}^n$ drawn from an unknown underlying distribution $p_{data}(x,y)$.

\textbf{Input Masking.}
A key ingredient of our approach is a masking function $\mathcal{G}$ that is applied offline (i.e., after full training). 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. This should encourage the model to investigate new features with the masked training set during finetuning. As for $\mathcal{G}$, we adopt the xGradCAM~\citep{selvaraju2017grad}, 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. Given an input image of size $H \times W \times C$, xGradCAM outputs a localization map $\mathcal{A}$ of size $H \times W \times 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. After acquiring the localization map, for each sample $(x_i, y_i)$, where $x_i \in X$ and $y_i \in Y$, we mask the locations with the most contribution as:

where $\mathcal{T}$ refers to a thresholding function by the threshold factor $\tau$ (i.e., $ \mathcal{T}=\mathbbm{1}_{\mathcal{A}_{x_i} \le \tau }$), and $\odot$ denotes element-wise multiplication. As the resolution of $\mathcal{A}$ is typically coarser than that of the input data, $\mathcal{T}(\mathcal{A}_{x_i})$ is up-sampled to the size of the input.

Procedurally, we first learn model $m_{\theta}^{\text{initial}}$ using original unmasked training data $\mathcal{D}^{\text{initial}}$. Then we use $m_{\theta}^{\text{initial}}$, $\mathcal{G}$ and $\mathcal{T}$ to create the masked set $\mathcal{D}^{\text{masked}}$. Finally, the fully trained predictor $m_{\theta}^\text{initial}$ is tuned using $\mathcal{D}^{\text{masked}}$ to obtain $m_{\theta}^{\text{final}}$. 
As for the masking step, any explainability approach can be applied (note that some may have more computational complexity, such as ScoreCAM~\citep{wang2020score}). We use xGradCAM~\citep{selvaraju2017grad} as it is fast and produces relatively denser heat-maps than other methods~\citep{srinivas2019full, selvaraju2017grad, wang2020score}. 

\subsection{Classification with Spurious Features} \label{sec:spurious}
\textbf{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. We construct a dataset such that a classifier would achieve poor accuracy by relying on spurious input features. 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. We use two test sets during testing: the original raw MNIST test set and a biased test set. 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). Figure~\ref{fig:mnist} demonstrates the performance of the ERM, RandMask, and our \textit{MaskTune} on both original and biased test sets. The RandMask method is similar to \textit{MaskTune}, but masks a \textit{randomly} chosen area of the input image. As shown in Figure~\ref{fig:mnist}, \textit{MaskTune} outperforms other methods by a large margin on both test sets. As demonstrated in Figure~\ref{fig:mnist} (right), \textit{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~\ref{app:exps}) and reported results for iterative version of ~\textit{MaskTune}.

The Waterbirds dataset~\citep{sagawa2019distributionally} was proposed to assess the degree to which models pick up spurious correlations in the training set. We discovered and fixed two \textit{issues} with the Waterbirds dataset: a) because the background images in the Places dataset~\citep{zhou2017places} may already 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~\citep{wah2011caltech}. For example, the label of an image may be ``landbird'', but the image contains both land and water birds. We manually removed such images from the dataset. b) Because the names of the species are similar, some land birds have been mislabeled as waterbirds which we corrected. The \textit{corrected} Waterbirds dataset can be found on \textit{MaskTune}'s GitHub page. After addressing the two issues, the ERM model's worst-group accuracy increased from 60\% which is reported in~\citep{sagawa2019distributionally} to 80.8$\pm$1.3\%. We repeated the group-DRO method and ERM experiments on the corrected waterbirds dataset and reported the results in Table~\ref{tab:waterbirds}. As demonstrated, our method (without any group supervision) achieves similar accuracy to the group-DRO that benefits from full supervision. In Figure~\ref{fig:celeba} (right), we visualized feature importance before and after applying our ~\textit{MaskTune} on the modified Waterbirds dataset.

Given input image $j$ and the original (\ensuremath{m^\text{initial}_\theta}) and finetuned (\ensuremath{m^\text{final}_\theta}) models, let their respective inference-time prediction probabilities for class $i$ be $P_{ij}^\text{init} $ and $P_{ij}^\text{final}$. Then \textit{MaskTune} does not abstain and declares the class $i$ as the predicted class, iff $P_{ij}^\text{init} \cdot P_{ij}^\text{final} > \gamma $, where $\gamma$ is a threshold that allows for controlling the target coverage. 
To find the proper $\gamma$ 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). 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.

\paragraph{More Aggressive Random Masking as a Baseline.} We compare \textit{MaskTune} to a another version of random masking with different window sizes on the MNIST dataset. For each image, we randomly selected a window of $\{2\times2, 3\times3, 4\times4, .... n \times n\}$ pixels where $n$ is the image size divided by $2$. As shown in Table \ref{tab:randmask2}, \textit{MaskTune} still outperforms the random masking approach.
{"value": "59.06", "type": "Result", "training data/set": "MNIST", "test data/set": "MNIST", "task": ["Classification with Spurious Features", "image classification"], "metric": "Accuracy", "experimental settings": {"Test data type": "biased"}, "model": "random masking", "model settings": {"xx": "yy"}} {"value": "91.02", "type": "Result", "training data/set": "MNIST", "test data/set": "MNIST", "task": ["Classification with Spurious Features", "image classification"], "metric": "Accuracy", "experimental settings": {"Test data type": "original"}, "model": "random masking", "model settings": {"xx": "yy"}} {"value": "98.26", "type": "Result", "training data/set": "MNIST", "test data/set": "MNIST", "task": ["Classification with Spurious Features", "image classification"], "metric": "Accuracy", "experimental settings": {"Test data type": "biased"}, "model": "MaskTune", "model settings": {"xx": "yy"}} {"value": "98.46", "type": "Result", "training data/set": "MNIST", "test data/set": "MNIST", "task": ["Classification with Spurious Features", "image classification"], "metric": "Accuracy", "experimental settings": {"Test data type": "original"}, "model": "MaskTune", "model settings": {"xx": "yy"}}