ASPIRE: Language-Guided Data Augmentation for Improving Robustness Against Spurious Correlations

TL;DR

ASPIRE tackles the problem of neural image classifiers leveraging spurious correlations by introducing a language-guided data augmentation pipeline that does not require minority-group labels. It identifies spurious cues via a six-step process that uses image captions and LLMs to extract foreground and background features, edits images to remove these cues, and personalizes a diffusion model to generate non-spurious in-domain augmentations. The augmented data are used to re-train classifiers, yielding consistent improvements in worst-group accuracy across four datasets and nine baselines, with particularly large gains on Hard ImageNet. This approach demonstrates a practical, annotation-free path to robust vision systems and complements existing robust-training methods, albeit with notable computational and captioning-quality considerations for future work.

Abstract

Neural image classifiers can often learn to make predictions by overly relying on non-predictive features that are spuriously correlated with the class labels in the training data. This leads to poor performance in real-world atypical scenarios where such features are absent. This paper presents ASPIRE (Language-guided Data Augmentation for SPurIous correlation REmoval), a simple yet effective solution for supplementing the training dataset with images without spurious features, for robust learning against spurious correlations via better generalization. ASPIRE, guided by language at various steps, can generate non-spurious images without requiring any group labeling or existing non-spurious images in the training set. Precisely, we employ LLMs to first extract foreground and background features from textual descriptions of an image, followed by advanced language-guided image editing to discover the features that are spuriously correlated with the class label. Finally, we personalize a text-to-image generation model using the edited images to generate diverse in-domain images without spurious features. ASPIRE is complementary to all prior robust training methods in literature, and we demonstrate its effectiveness across 4 datasets and 9 baselines and show that ASPIRE improves the worst-group classification accuracy of prior methods by 1% - 38%. We also contribute a novel test set for the challenging Hard ImageNet dataset.

Figures (8)

• Figure 1: Overview of ASPIRE. Given a training dataset, ASPIRE automatically detects non-predictive spuriously correlated features for each class (e.g., indoor background for small dogs) and generates synthetic images without them (small dogs in an outdoor background). These images can then be added to the train set to learn a more robust image classifier.
• Figure 2: Illustration of ASPIRE: ASPIRE follows a 6-step process to improve the robustness against spurious correlations. 1 We first train a base classifier $\mathcal{E}$ using ERM on the entire training set and extract images with features that are spuriously correlated to construct $\mathcal{D}_{hold}$. 2 We caption each image in $\mathcal{D}_{hold}$. 3 We feed the caption to a LLM and extract the foreground objects and background for each image. 4a We remove one foreground object at a time and predict the class of the edited image $\mathcal{E}$. If $\mathcal{E}$ predicts incorrectly, we consider the object as a plausible spurious correlation learned by $\mathcal{E}$ for that class. 4b We edit the image to change its original background with an alternative background suggested by the LLM and follow the process to similar to 4a. 5 We personalize a text-to-image diffusion model using edited images from the previous step for the top-k unique items leading to the highest number of wrong predictions. 6 We re-train $\mathcal{E}$ using the generated augmentations to obtain $\hat{\mathcal{E}}$.
• Figure 3: GradCAM visualizations of the features used by the standard ERM model trained with and w/o ASPIRE augmentations on the Hard ImageNet dataset (Balance beam top and Volleyball bottom). As clearly visible, when trained with ASPIRE augmentations, the model tends to focus better on core features than spurious ones (more in Appendix \ref{['sec:gradcam']}).
• Figure 4: Examples of Original Images, Edited Images from the ASPIRE pipeline and Generated Augmentations. To the left of the Generated Augmentations, we also mention the top-k spurious correlations discovered by ASPIRE for the particular class. ASPIRE generates diverse augmentations with the desired non-spurious features that can be used to train robust models.
• Figure 5: Images illustrating cases of ASPIRE failures.