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

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.