You Only Need a Good Embeddings Extractor to Fix Spurious Correlations

TL;DR

This paper demonstrates that using embeddings from large pre-trained vision models with a simple linear classifier can effectively mitigate spurious correlations in datasets like Waterbirds, achieving comparable or better robustness than specialized training methods.

Contribution

It shows that a straightforward approach using pre-trained embeddings and linear classifiers can surpass existing methods that require subgroup labels for mitigating spurious correlations.

Findings

Pre-trained vision transformers outperform CNNs in robustness.

Larger pre-training datasets improve worst-group accuracy.

Simple linear classifiers on embeddings can match advanced robustness methods.

Abstract

Spurious correlations in training data often lead to robustness issues since models learn to use them as shortcuts. For example, when predicting whether an object is a cow, a model might learn to rely on its green background, so it would do poorly on a cow on a sandy background. A standard dataset for measuring state-of-the-art on methods mitigating this problem is Waterbirds. The best method (Group Distributionally Robust Optimization - GroupDRO) currently achieves 89\% worst group accuracy and standard training from scratch on raw images only gets 72\%. GroupDRO requires training a model in an end-to-end manner with subgroup labels. In this paper, we show that we can achieve up to 90\% accuracy without using any sub-group information in the training set by simply using embeddings from a large pre-trained vision model extractor and training a linear classifier on top of it. With experiments on a wide range of pre-trained models and pre-training datasets, we show that the capacity of the pre-training model and the size of the pre-training dataset matters. Our experiments reveal that high capacity vision transformers perform better compared to high capacity convolutional neural networks, and larger pre-training dataset leads to better worst-group accuracy on the spurious correlation dataset.