Are Vision Transformers Robust to Spurious Correlations?

Abstract

Deep neural networks may be susceptible to learning spurious correlations that hold on average but not in atypical test samples. As with the recent emergence of vision transformer (ViT) models, it remains unexplored how spurious correlations are manifested in such architectures. In this paper, we systematically investigate the robustness of different transformer architectures to spurious correlations on three challenging benchmark datasets. Our study reveals that for transformers, larger models and more pre-training data significantly improve robustness to spurious correlations. Key to their success is the ability to generalize better from the examples where spurious correlations do not hold. Further, we perform extensive ablations and experiments to understand the role of the self-attention mechanism in providing robustness under spuriously correlated environments. We hope that our work will inspire future research on further understanding the robustness of ViT models to spurious correlations.

Appendices

Appendix A Implementation Details

• 1. Transformers. For ViT models, we obtain the pre-trained checkpoints from the timm library.¹ For downstream fine-tuning on the Waterbirds and CelebA dataset, we scale up the resolution to 384 $\times$ 384 by adopting 2D interpolation of the pre-trained position embeddings proposed in Dosovitskiy et al. (2021). Note, for CMNIST we keep the resolution as $224 \times 224$ during fine-tuning. We fine-tune models using SGD with a momentum of 0.9 with an initial learning rate of 3e-2. As described in (Steiner et al., 2021), we use a fixed batch size of 512, gradient clipping at global norm 1, and a cosine decay learning rate schedule with a linear warmup.

• 2. BiT. We obtain the pre-trained checkpoints from the official repository.² For downstream fine-tuning, we use SGD with an initial learning rate of 0.003, momentum 0.9, and batch size 512. We fine-tune models with various capacities for 500 steps, including BiT-M-R50x1, BiT-M-R50x3, and BiT-M-R101x1.

• 3. Data Augmentation Schemes. For applying different data augmentations and regularization approaches in Sect. 1, we use the helper functions provided with timm library. We report different hyper-parameters used in Table 7.

Table 7 Different hyper-parameters used for data augmentation approaches in Sect. 1

Appendix B Extension: How Does the Size of Pre-training Dataset Affect Robustness to Spurious Correlations?

In this section, to further validate our findings on the importance of large-scale pre-training dataset, we show results on CelebA (Liu et al., 2015) dataset. We report our findings in Table 8. We also observe a similar trend for this setup that larger model capacity and more pre-training data yield a significant improvement in worst-group accuracy. Further, when pre-trained on a relatively smaller dataset such as ImageNet-1k, the performance of both transformer and CNN models are poor as compared to the ImageNet-21k counterpart.

Also, compared to BiT models, the robustness of transformer models benefits more with a large pre-training dataset. For example, compared to ImageNet-1k, fine-tuning DeiT-III-Base pre-trained on ImageNet-21k improves the worst-group accuracy by 6.5%. On the other hand, for BiT models, fine-tuning with a larger pre-trained dataset yields marginal improvement. Specifically, BiT-M-R50x3 only improves the worst-group accuracy by 1.5% with ImageNet-21k.

Table 8 Investigating the effect of large-scale pre-training on model robustness to spurious correlations when finetuned on CelebA (Liu et al., 2015)

Appendix C Spurious Out-of-Distribution Detection

Table 9 Spurious OOD evaluation.

In this section, we study the performance of ViT models in out-of-distribution setting. Introduced in Ming et al. (2022), spurious out-of-distribution (OOD) data is defined as samples that do not contain the invariant features $\textbf{z}^\textrm{inv}$ essential for accurate classification, but contain the spurious features $\textbf{z}^{e}$. Hence, these samples are denoted as $\textbf{x}_\textrm{ood} = \rho (\textbf{z}^{\bar{y}},\textbf{z}^e)$ where $\bar{y}$ is an out-of-class label, such that $\bar{y} \not \in \mathcal {Y}$. In the problem of waterbird vs landbird classification, an image of a person standing in forest would be an example of spurious OOD, since it contains different semantic class person $\not \in \{\texttt {waterbird}, \texttt {landbird}\}$, yet has the environmental features of land background. A non-robust model relying on the background feature may classify such OOD data as an in-distribution class with high confidence. Hence, we aim to understand if self-attention based ViT models can mitigate this problem and if so, to what extent.

To investigate the performance of different models against spurious OOD examples, we use the setup introduced in Ming et al. (2022). Specifically, for Waterbirds (Sagawa et al., 2020) we test on subset of images of land and water sampled from the Places dataset (Zhou et al., 2017). Considering, CelebA (Liu et al., 2015) as in-distribution, our test suite consists of images of bald male as spurious OOD, since they contain environmental features (gender) without invariant features (hair). For CMNIST, the in-distribution data contains digits $\mathcal {Y}$ = $\{0,1\}$ and the background colors, $\mathcal {E}$ = {red, green, purple, pink}. We use digits {5, 6, 7, 8, 9} with background color red and green as test OOD samples. We report our findings in Table 9. Clearly, ViT models achieve better OOD evaluation metrics as compared to BiTs. Specifically, ViT-B/16 achieves $+{32}\%$ higher AUROC than BiT-M-R50x3, considering Waterbirds (Sagawa et al., 2020) as in-distribution.

Appendix D Extension: Color Spurious Correlation

To further validate our findings beyond natural background and gender as spurious (i.e. environmental) features, we provide additional experimental results with the ColorMNIST dataset, where the digits are superimposed on colored backgrounds. Specifically, it contains a spurious correlation between the target label and the background color. Similar to the setup in Ming et al. (2022), we fix the classes $\mathcal {Y}$ = $\{0,1\}$ and the background colors, $\mathcal {E}$ = {red, green, purple, pink}. For this study, label $y=0$ is spuriously correlated with background color $\{\texttt {red}, \texttt {purple}\}$, and similarly, label $y=1$ has spurious associations with background color $\{\texttt {green}, \texttt {pink}\}$. Formally, we have ${\mathbb {P}}(e = \texttt {red} \mid y = 0) = {\mathbb {P}}(e = \texttt {purple} \mid y = 0) = {\mathbb {P}}(e = \texttt {green} \mid y = 1) = {\mathbb {P}}(e = \texttt {pink} \mid y = 1) = 0.45$ and ${\mathbb {P}}(e = \texttt {green} \mid y = 0) = {\mathbb {P}}(e = \texttt {pink} \mid y = 0) = {\mathbb {P}}(e = \texttt {red} \mid y = 1) = {\mathbb {P}}(e = \texttt {purple} \mid y = 1) = 0.05$. Note that, while fine-tuning the models, we fix the foreground color of digits as $\texttt {white}$.

Results and insights on robustness performance. We compare model predictions on samples with the same class label but different background & foreground colors. Given a data point ($\textbf{x}_i, y_i$), we modify the background and foreground color of $\textbf{x}_i$ randomly to generate a new test image $\bar{\textbf{x}}_i$ with the constraint of having the same semantic label. In evaluation, the background color is chosen uniform-randomly from the set of colors: {#ecf02b, #f06007, #0ff5f1, #573115, #857d0f, #015c24, #ab0067, #fbb7fa, #d1ed95, #0026ff} and the foreground color is selected randomly from the set $\{\texttt {black},\texttt {white}\}$. For evaluation, we form a dataset consisting of 2100 samples. The results reported are averaged over 50 random runs. Figure 6 depicts the distribution of training samples in the CMNIST dataset (left) and some representative examples after transformation (right).

We report our findings in Fig. 7. Our operating hypothesis is that a robust model should predict the same class label $\hat{f}(\textbf{x}_i)$ and ${\hat{f}}(\bar{\textbf{x}}_i)$ for a given pair $(\textbf{x}_i, \bar{\textbf{x}}_i)$, as they share exactly the same target label (i.e., the invariant feature is approximately the same). We can observe from Fig. 7 that the best model ViT-B/16 obtains consistent predictions for 100% of image pairs. After extensive experimentation over all combinations, we find that setting the foreground color as $\texttt {black}$ and the background as $\texttt {white}$ caused the models to be most vulnerable. We see a significant decline in model consistency when the foreground color is set as $\texttt {black}$ and the background as $\texttt {white}$ (indicated as BW) as compared to random setup.

Fig. 6

CMNIST. Distribution of training samples in CMNIST dataset (left) and few representative examples after transformation (right) as defined in Section D

Fig. 7

Consistency Measure. Evaluation results quantifying consistency for models of different architectures and varying capacities. We indicate the setup when the foreground color is set as $\texttt {black}$ and the background as $\texttt {white}$ using BW (right). Random represents setting both the foreground and background color randomly (left)

Appendix E Visualization

1.1 E.1 Attention Map

In Fig. 8, we visualize attention maps obtained from ViT-B/16 model for some samples in Waterbirds (Sagawa et al., 2020) and CMNIST dataset. We use Attention Rollout (Abnar & Zuidema, 2020) to obtain the attention matrix. We can observe that the model successfully attends spatial locations representing invariant features while making predictions.

Fig. 8

Attention map. Visual illustration of attention map obtained from ViT-B/16 model for few representative images

1.2 E.2 The Attention Matrix of CMNIST

In the main text, we provide visualizations in which each image patch, irrespective of its spatial location, provides maximum attention to the patches representing essential cues for accurately identifying the foreground object. In Fig. 9, we show visualizations for ViT-B/16 fine-tuned on the CMNIST dataset to further validate our findings.

Fig. 9

Visualization of the top N patches receiving the highest attention (marked in black) for ViT-B/16 fine-tuned on CMNIST. Investigating the attention matrix, we find that all image patches-irrespective of spatial location-provide maximum attention to the patches representing essential cues

Table 10 Investigating the effect of large scale pre-training on model robustness to spurious correlations when finetuned on Waterbirds (Sagawa et al., 2020) dataset

Table 11 Investigating the effect of training models on Waterbirds (Sagawa et al., 2020) dataset from scratch without any pre-training

Appendix F Extension: Pattern in Attention Matrix

In this section, we provide visualizations for top N patches receiving the highest attention values for ViT-Ti/16 (Fig. 10) fine-tuned on Waterbirds dataset (Sagawa et al., 2020) on various test images.

Takeaways. Figure 10 shows the top N patches receiving the highest attention values for ViT-Ti/16 model fine-tuned on the Waterbirds dataset. We observe: (1) The model correctly attends to patches responsible for accurate classification in images belonging to the majority groups, i.e, waterbird on water background and landbird on land background. (2) For images belonging to minority groups (3rd and 4th row in Fig. 10) such as waterbird on land background, the model provides maximum attention to the environmental features, exhibiting lack of robustness to spurious correlations.

Fig. 10

Visualization of the top N patches receiving the highest attention (marked in red) for ViT-Ti/16 model finetuned on Waterbirds (Sagawa et al., 2020) dataset (Color figure online)

Appendix G Extension: Results Highlighting the Importance of Pre-training Dataset on Additional Architectures

In this section, to further validate our findings on the importance of large-scale pre-training dataset we show results for additional transformer architectures: RVT (Mao et al., 2022) and PVT (Wang et al., 2021).

Table 12 Ablations on using different configurations for fine-tuning DeiT-III-Small model on Waterbirds dataset (Sagawa et al., 2020)

Table 13 Ablations on using different configurations for fine-tuning BiT-R50x1 model on Waterbirds dataset (Sagawa et al., 2020)

Takeaways: We report the results of finetuning different model architectures on Waterbirds (Sagawa et al., 2020) dataset in Table 10. Note, RVT architecture is specifically modified to provide additional robustness against adversarial attacks, common corruptions, and out-of-distribution inputs. Despite these modifications, RVT models when pre-trained on ImageNet-1k perform poorly when fine-tuned on Waterbirds indicating a lack of robustness to spurious correlations. However, we do notice that RVT displays stronger robustness to spurious correlation than ViT when pre-trained on ImageNet-1k.

Appendix H Experiments on training models from scratch

In Table 11, we train ViT and BiT models from scratch on Waterbirds(Sagawa et al., 2020) dataset. We observe that without any pre-training, both ViT and BiT models severely overfit the training dataset to attain 100% training accuracy and significantly fail on test samples. This observation indicates that without pre-training, both transformers and CNNs have a high propensity of memorizing training samples (along with their inherent bias).

Fig. 11

GradCAM visualizations for few images sampled from Waterbirds dataset (Sagawa et al., 2020). The model used is BiT-S-R50x1. We use GradCAM (Selvaraju et al., 2016) to visualize the “salient” observed features used to classify images

Appendix I How Does Different Training Configuration Affect Robustness to Spurious Correlations?

In this section, we aim to disentangle the effect of different training configurations on robustness to spurious correlations for DeiT-III (Touvron et al., 2022) models. Unlike ViT (Dosovitskiy et al., 2021), DeiT-III models are pre-trained using strong data augmentations. Further, for ImageNet (Deng et al., 2009) training, Touvron et al. (2021) have shown the positive impact of different data augmentation and regularization approaches on model accuracy. Motivated by the findings, we aim to verify if applying different training schemes during fine-tuning can further improve model robustness to spurious correlations.

Table 12 reports results for fine-tuning on Waterbirds(Sagawa et al., 2020) using different configurations.³ For this experiment, we use the DeiT-III-Small architecture pre-trained on ImageNet-21k. The top row in Table 12 corresponds to the standard fine-tuning setting we use throughout the paper for consistency and fair evaluation between different architectures. Considering data augmentation schemes, we observe that Mixup (Zhang et al., 2018), CutMix (Yun et al., 2019), and Rand Augment (Cubuk et al., 2020) provide the most significant improvement in model robustness. Specifically, using Rand Augment (Cubuk et al., 2020) during fine-tuning leads to 4.98% improvement in worst-group accuracy. Surprisingly, we also observe that simultaneously applying too many data augmentation schemes can instead hamper the worst-group accuracy. Among regularization schemes, we observe that reducing the weight decay penalty can also provide some improvement in worst-group accuracy.

Next, we investigate the impact of different training schemes on the robustness of convolutional nets to spurious correlations. Precisely, we fine-tune the BiT-R50x1 model on the Waterbirds dataset using different data augmentation approaches. We report results in Table 13. We observe that, irrespective of the pre-training dataset, using Auto Augment (Cubuk et al., 2019) and Rand Augment (Cubuk et al., 2020) improves model robustness to spurious correlations.

Appendix J GradCAM Visualizations

In Fig. 11, we show GradCAM (Selvaraju et al., 2016) visualizations of a few representative examples from the Waterbirds (Sagawa et al., 2020) dataset for a BiT-S-R50x1 model. For each input image, we show saliency maps where warmer colors depict higher saliency. We observe that the model correctly attends to features responsible for accurate classification in images belonging to the majority groups, i.e., waterbird on water background and landbird on land background. However, for the image belonging to the minority group, i.e., waterbird on land background, the model provides maximum attention to the environmental features, exhibiting a lack of robustness to spurious correlations. This visualization directly corroborates our empirical findings in Table 3.

Appendix K Software and Hardware

We run all experiments with Python 3.7.4 and PyTorch 1.9.0 using Nvidia Quadro RTX 5000 GPU.