Correct-N-Contrast: A Contrastive Approach for Improving Robustness to Spurious Correlations

TL;DR

The paper tackles robustness to spurious correlations by proposing Correct-n-Contrast (CnC), a two-stage method that uses contrastive learning to align representations by class while reducing reliance on spurious attributes. Stage 1 extracts pseudo group labels from an ERM model, and Stage 2 employs a supervised contrastive objective with a novel sampling scheme to pull together same-class samples across different spurious attributes and push apart different classes with the same spurious signal. The authors establish a theoretical link between alignment loss and worst-group performance, demonstrate substantial improvements on multiple benchmarks without group labels, and confirm via representation metrics and ablations that CnC learns more class-centered, spurious-attribute-robust features. Overall, CnC achieves state-of-the-art worst-group accuracy on several spuriously correlated datasets and remains competitive with oracle-group methods.

Abstract

Spurious correlations pose a major challenge for robust machine learning. Models trained with empirical risk minimization (ERM) may learn to rely on correlations between class labels and spurious attributes, leading to poor performance on data groups without these correlations. This is particularly challenging to address when spurious attribute labels are unavailable. To improve worst-group performance on spuriously correlated data without training attribute labels, we propose Correct-N-Contrast (CNC), a contrastive approach to directly learn representations robust to spurious correlations. As ERM models can be good spurious attribute predictors, CNC works by (1) using a trained ERM model's outputs to identify samples with the same class but dissimilar spurious features, and (2) training a robust model with contrastive learning to learn similar representations for same-class samples. To support CNC, we introduce new connections between worst-group error and a representation alignment loss that CNC aims to minimize. We empirically observe that worst-group error closely tracks with alignment loss, and prove that the alignment loss over a class helps upper-bound the class's worst-group vs. average error gap. On popular benchmarks, CNC reduces alignment loss drastically, and achieves state-of-the-art worst-group accuracy by 3.6% average absolute lift. CNC is also competitive with oracle methods that require group labels.

Figures (15)

• Figure 1: (1) ERM-trained models classify by spurious features, shown via GradCAM selvaraju2017grad. (2) CnC learns similar representations for same-class samples with different ERM predictions to ignore spurious attributes and classify samples correctly.
• Figure 2: ERM-learned representation UMAPs, trained on spuriously uncorrelated ($p_\text{corr} = 0.2$) vs correlated ($p_\text{corr} = 0.95$) data.
• Figure 3: Accuracy and representation metrics from ERM models trained on increasingly spuriously correlated Colored MNIST. Lower worst-group accuracy (Fig. 3a) corresponds to both higher alignment loss (Fig. 3b) and $\hat{I}(Y; Z) < \hat{I}(A; Z)$ (Fig. 3c, Fig. 3d).
• Figure 4: Higher worst-group accuracy with Jtt (versus Fig. \ref{['fig:erm_ablation_cmnist']}a) coincides with keeping $\hat{I}(Y; Z) \gg \hat{I}(A; Z)$.
• Figure 5: Model alignment loss (a) and mutual information (b, c) after training with ERM, Jtt, and CnC. CnC most effectively reduces spurious attribute dependence, and obtains smaller gaps for per-class worst-group versus average error (d), as supported by Theorem \ref{['prop_close']}.

Theorems & Definitions (4)

• Theorem 3.1
• proof : Proof of Theorem \ref{['prop_close']}
• Corollary B.1: Extension of Theorem \ref{['prop_close']} to compare across different classes
• proof