Identifying and Disentangling Spurious Features in Pretrained Image Representations

In many classification datasets, some features are predictive of the label but are not causally related. It is often said that these features are spuriously correlated with the label, as their correlation might not hold for data collected in another environment. For example, suppose we collect typical images of cows and camels and form a binary classification task. In that case, we will find that the background is correlated with the label, as cows are often photographed in barns or green pastures, while camels are often photographed in deserts (Beery et al., 2018). However, this correlation will be spurious as the background information is not causally related to the label, and we can easily make another dataset of cows and camels in which this correlation does not hold.

It is well-established that neural networks are susceptible to spurious correlations (Torralba & Efros, 2011; Ribeiro et al., 2016; Gururangan et al., 2018; Zech et al., 2018; McCoy et al., 2019; Geirhos et al., 2019, 2020; Xiao et al., 2021). In such cases, neural networks learn representations that capture spurious features and make predictions that employ them. Many approaches have been proposed for learning representations that do not capture spurious features (Muandet et al., 2013; Sun & Saenko, 2016; Ganin et al., 2016; Wang et al., 2019b, a; Li et al., 2018; Arjovsky et al., 2019; Zhao et al., 2020; Lu et al., 2022). Some methods are tailored against specific spurious correlations (e.g., texture); some require specifying a categorical spurious feature, while others require data collected from multiple labeled environments. Nevertheless, to our best knowledge, none of such representation learning methods consistently outperform standard empirical risk minimization (Gulrajani & Lopez-Paz, 2021; Koh et al., 2021). This is partly because spurious features are often easier to learn and get learned early in training (Shah et al., 2020; Nam et al., 2020; Hermann & Lampinen, 2020; Pezeshki et al., 2021).

Besides the unsatisfactory results, the approach mentioned above also goes against one of the main techniques of deep learning – using pretrained representations instead of learning from scratch. Recently, a few works indicated a large potential in fixing pretrained representations and focusing on training a linear classifier on top of it that does not rely on spurious correlations. In particular, Galstyan et al. (2022) find that a significant contribution to the out-of-domain generalization error comes from the classification head and call for designing better methods of training the classification head. Menon et al. (2021) propose to retrain the classification head on training data with down-sampled majority groups. Kirichenko et al. (2023); Izmailov et al. (2022); and Shi et al. (2023) find that after training on data with spurious correlations, keeping the representations fixed and retraining the classification head on small unbiased data gives state-of-the-art results. When no information about spurious features is available, Mehta et al. (2022) show that one can still get good results by using embeddings from a large pretrained vision model. Interestingly, representations learned by a vision transformer (Dosovitskiy et al., 2021) seem to lead to more robust classification heads (Ghosal et al., 2022). Overall, these findings indicate that more research is needed to understand better how spurious features are represented and design better methods of training classification heads on representations that capture spurious features.

We consider the Waterbirds dataset (Sagawa* et al., 2020), which is landbird vs waterbird image classification task where the background is spuriously correlated with the label. Namely, most landbird images have land in their background, while most waterbird images have water in their background. We consider fixed pretrained representations learned through supervised or self-supervised learning. We investigate whether one can remove the spurious features from the representations in two settings. In the former (and more prevalent setting), one has access to the value of the binary spurious feature. In the latter, we also have access to per-example image masks indicating which parts of images correspond the spurious feature.

Interestingly, even with full knowledge of the spurious feature, it is not straightforward to remove it. While we find that representations are axis-aligned to a certain degree, the extent of alignment is not enough to remove spurious features by removing individual representation coordinates. Since both the spurious feature and the label can be predicted well from the representation with a linear layer, we hypothesize that the entanglement of core and spurious features is linear and can be reversed with a linear transformation. For this we propose a linear autoencoder to split the representation into three parts corresponding to the class label, the spurious feature, and other features not related to the former two but required for reconstruction. Importantly, in contrast to existing approaches, we do not enforce independence of the first and second parts on a biased training set. Instead, we enforce independence on an upsampled variant of the training set.

We find that a linear classifier trained on the core features of the encoding performs better than the standard approach but does not reach the performance of a classifier trained on an unbiased set. We demonstrate that this gap can be closed by performing additional feature selection within the core features.

Waterbirds dataset. We consider the Waterbirds dataset (Sagawa et al., 2020), which is a benchmark dataset designed to measure the effect of debiasing spurious correlations. The dataset consists of bird photographs from the CUB (Welinder et al., 2010) dataset combined with image backgrounds from the Places (Zhou et al., 2018) dataset. It has 4,795 training examples, 1,199 val examples, and 5,794 testing examples. The task is to classify birds as waterbirds or landbirds while ignoring the background.

Bird masks. For every picture of the training set we use Detic (Zhou et al., 2022) to segment the bird in the image. The binary mask of the bird in the picture is denoted by $m$. We visually evaluate the quality of the masks and they are close to ideal. The most common error is that in rare cases there are additional birds in the background which are included in the mask.

Pretrained representations. We use three pretrained models in our experiments: ImageNet-pretrained ResNet-50, SWAG-pretrained and ImageNet-finetuned RegNetY, and a self-supervised ViT-B/14 from DINOv2. These models produce $d$-dimensional representations for each, where $d=2048$, $d=7392$, $d=786$ for the three models respectively. Mehta et al. (2022) shows that the better models produce better worst group accuracy (WGA). Throughout this paper we use $z$ to denote representations.

Attributing neurons. Whether or not individual neurons can be attributed to specific image regions is not clear. There is some evidence against it and some evidence in support of it (Elhage et al., 2022). Certainly some layers and operations in deep learning give preference to the standard basis. These include element-wise activations functions, batch normalization, dropout, etc. We use Captum (Kokhlikyan et al., 2020) to find this attribution. Captum has implementations of several attribution algorithms. In our preliminary experiments, we found that the results with the Integrated Gradients method (Sundararajan et al., 2017) are good enough.

For a given image $x_{j}\in\mathbb{R}^{224\times 224}$, with foreground mask $m_{j}\in\left\{0,1\right\}^{224\times 224}$, and representation $z_{j}\in\mathbb{R}^{d}$, we use Captum to compute attribution heatmap $a^{i}_{j}\in\mathbb{R}^{224\times 224}$ of the neuron $i$-th neuron over input pixels. We define spuriousness $s(x_{j})\in\mathbb{R}^{d}$ the following way:

where the sums are over pixels. The first term is the average attribution on the foreground pixels and the second term is the average attribution on the background pixels. Note that the attribution scores given by Captum can be both negative and positive, indicating the direction of the impact on the individual neuron. We use absolute values of the attribution scores as we are only interested in the magnitude of the impact. The spuriousness of the specific neuron $z_{i}$ is defined as the average spuriousness over the training set $X$:

Linear models. We train a linear classifier on a frozen representation using the scikit-learn package. We use L-BFGS optimizer and disable regularization. We report overall accuracy and worst group accuracy on the test set, and mostly focus on optimizing the latter. We note that SGD-based methods to learn a linear classifier can discover drastically different solutions, including solutions with higher worst group accuracy. This is especially true when early stopping is used. All linear models used in this paper use L-BFGS, and the analysis of SGD-based optimization is left for future work.

Upper bound. We compute an upper bound on worst group accuracy of L-BFGS-trained linear models by splitting the test set into five equal parts of 1000 samples, and perform five-fold cross-validation. This way we ensure that each of the five models is trained on a subset of the same size as the training set, but with equal number of images with water backgrounds and land backgrounds. The average of worst group accuracies of the five models is reported as an upper bound.

Statistical significance. For selected experiments we performed bootstrapping on the training set to estimate the variability of the models with respect to changes in the training set. We resampled with repetitions five versions of the training set of the same size, and repeated the experiments on the five versions. We report the mean and standard deviation of the five metrics.

We start our investigations by studying whether one can remove spurious features by removing individual neurons from the pretrained representations.

Removing features that look mostly to background improves worst group accuracy. We sort the $d$ neurons according to $s_{i}$ and consider keeping only the top $N$ neurons for the linear models. These models are denoted by $\text{Captum}^{N}(z,m)$. Figure 1 shows the results. In case of ResNet-50, a linear model on the top $N=50$ neurons significantly improves the baseline (see also Table 1). The improvement is seen with up to $N=260$ neurons. This confirms the hypothesis that there are many spurious features that harm the worst group accuracy of linear models, and $s_{i}$ can be used as a measure to identify them. A similar effect is observed with RegNetY, there is improvement for at least up to $N=1000$.

We could not detect this phenomenon in case of DINOv2. Keeping top neurons in terms of $s_{i}$ worsens the metrics. Most likely this means that the individual neurons are neither pure spurious nor pure non-spurious. A supporting evidence is that the variance of $s_{i}$ scores across the neurons is smaller than in case of ResNet-50.

Pretrained representations are relatively axis-aligned. To verify whether the directions in the representation space responsible for spurious features are aligned with the axes, we sample a random rotation matrix, apply it to the representations $z$, and recalculate $s_{i}$ for them. Then we pick the top neurons from the rotated space and train new linear models. These models are denoted by $\text{Captum}^{N}(\text{Rot}(z),m)$. As seen in Figure 1, the scores are significantly worse in ResNet-50 (for $N\leq 250$) and RegNetY (for $N\leq 150$). This implies that the spuriousness directions are aligned with axes for these two representations. The difference is much smaller in DINOv2, which is expected, as there were no distinctive spurious neurons even before the rotation.

Designing a group-aware autoencoder. As seen in the previous section, there exist spurious coordinates in ResNet-50 and RegNetY representations. We also showed that distinctive spuriousness is lost when we apply a random rotation matrix. This raises a question whether there exists another linear transformation that will make spurious features even more axis aligned, i.e. there will be new neurons that more specifically capture the spurious features. In other words, we are looking for ways to disentangle spurious and core features with a linear transformation.

We design a simple autoencoder where the linear encoder maps input $z$ to three vectors: $z_{y}$, $z_{c}$ and $z_{n}$. We force $z_{y}$ and $z_{c}$ to contain information about the label and the background, respectively. We do this by adding another linear layer on top of $z_{y}$ and $z_{c}$ that predict $\hat{y}=W_{y}z_{y}$ and $\hat{c}=W_{c}z_{c}$ which are supervised by the corresponding signals. The linear decoder takes the concatenation of $z_{y}$, $z_{c}$ and $z_{n}$ and reconstructs $\hat{z}$ which should be close to the original $z$. We assume that $z_{n}$ will store the rest of the information in $z$ that is not relevant for predicting either the label or the background. Note that we need group-level information to train this autoencoder, but we do not need the masks of the birds. Following (Jaiswal et al., 2019), we add an additional regularization term that minimizes mutual information between $z_{y}$ and $z_{c}$. The final loss function is the following:

where $\text{ce}(\cdot,\cdot)$ is the cross-entropy loss, and HSIC denotes the Hilbert-Schmidt Independence Criterion (Gretton et al., 2005).

We train the autoencoder on the upsampled version of the training set so that each group is represented equally. This also justifies the minimization of the mutual information between $z_{y}$ and $z_{c}$, as they are correlated in the original training distribution. The sum of the dimensions of the three vectors $z_{y}$, $z_{c}$ and $z_{n}$ matches $d$ for each backbone. The linear model trained on top of the concatenation of the three vectors is denoted by $\text{GwAE}(z,g)$, where GwAE denotes the linear encoder of the group-wise trained autoencoder, and $g$ refers to the group information required for the training. $\text{GwAE}_{y}(z,g)$ denotes only the $z_{y}$ part of the encoder’s output. Note that the linear classifier trained on top of the autoencoder still belongs to the space of linear classifiers.

We expect the linear models trained on $z$ and on the full $\text{GwAE}(z,g)$ to perform similarly. We surprisingly see that this is not the case with ResNet-50. $\text{GwAE}(z,g)$ is better by 6 percentage points. One explanation is that the autoencoder is not ideal, some information is lost by the encoder, and luckily the lost information contains some of the spurious features. We leave a deeper analysis for the future work.

Label-aware part has mostly good features. For all backbones we see that $z_{y}$ gathers core features and the linear models trained on them have significantly better worst group accuracy. In case of ResNet-50 and DINOv2, the results are even better than the ones by Captum, which means that the autoencoder managed to isolate core features much better than it was possible by simply removing neurons in the original $z$. In case of RegNetY, we see that the linear model on $z_{y}$ performs as good as many models trained on Captum-filtered neurons. This means that RegNetY features were already disentangled.

Furthermore, we apply Captum on top of the $z_{y}$ neurons to see whether we can still identify and remove spurious features in $z_{y}$ (the orange plot in Figure 1). This was successful only in case of ResNet-50.

PCA helps. To separate the core features, the linear encoder needs to shrink some directions (for example those corresponding to spurious features). For this reason, one can hypothesize that most of the variance in $z_{y}$ will be along the core features. This motivates applying principal component analysis (PCA) on $z_{y}$ to further remove non-core features. Unlike Captum, PCA does not require additional information from the data. We find that, indeed, training linear models on the $N$ principal components of $z_{y}$ still improves the worst group accuracy for ResNet and RegNetY. We note that applying PCA directly on $z$ does not help to identify core features: the principal components of the original space usually contain spurious features.

With carefully designed experiments we have shown that pretrained image representations contain spurious features that can be identified and removed to improve worst group accuracy of the linear models. In most representations it is also possible to disentangle spurious features and further improve the performance. In future work we plan experiments with more backbones and datasets to see how well these findings generalize to other settings. We also note that it is possible to find linear models with better worst group accuracy if we use stochastic gradient descent with early stopping. This analysis is also left for future work.