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

1 INTRODUCTION . For many tasks , deep neural networks are negatively affected by spurious correlations—dependencies between observed features and class labels that only hold for certain groups of the data . For example , consider classifying images of cows or camels , where 90 % of cow images depict grassy backgrounds . A model may learn to predict the “ cow ” class based on the background , and incorrectly classify cow images with non-grass backgrounds as camels ( Ribeiro et al. , 2016 ; Beery et al. , 2018 ; Kaufman et al. , 2012 ) . This illustrates a widespread issue where neural networks can achieve low test error on certain groups , yet high error on others ( Blodgett et al. , 2016 ; Buolamwini & Gebru , 2018 ; Hashimoto et al. , 2018 ; Sagawa et al. , 2019 ) . Prior works have shown that this problem is increasingly aggravated as the correlations between class labels and spurious features become stronger ( Sagawa et al. , 2020 ) and easier to learn ( Arpit et al. , 2017 ; Hermann & Lampinen , 2020 ) . Since spurious correlations arise in many settings , we wish to design robust methods that perform well on all groups . How can we obtain neural networks robust to spurious correlations ? If group-defining information ( i.e . spurious attributes ) is known , a common solution is to minimize the worst-group loss , e.g. , with group DRO ( GDRO ) ( Sagawa et al. , 2019 ) . However , such information may be expensive to collect , and we may not know the spurious attributes a priori in a given data set ( Oakden-Rayner et al. , 2020 ) . When group information is unavailable , prior works typically take a two-stage approach . They first train an ERM model , and then use this model to infer groups and train a more robust model . For example , Sohoni et al . ( 2020 ) find that ERM models still learn group-specific features when trained to predict class labels . After first training an ERM model , they infer groups by clustering the ERM Under review as a conference paper at ICLR 2022 Landbird , Land BG ✅ ✅ ✅ ✅ ERM Ours Landbird , Water BG Landbird , Land BG Waterbird , Land BG Waterbird , Water BG Landbird , Water BG Landbird , Land BG Waterbird , Land BG Waterbird , Water BG ❎ ✅❎✅ ERM ✓ ✗ GradCAM Correct ? Correct ? “ Waterbird ” “ Landbird ” “ Waterbird ” “ Waterbird ” “ Landbird ” “ Waterbird ” ✗ Sample contrastive batches “ Landbird ” “ Waterbird ” “ Waterbird ” model ’ s representations , and train a new model with GDRO using these inferred groups . Creager et al . ( 2021 ) identify groups under which an initial trained ERM model would maximally violate the invariant risk minimization ( IRM ) objective ( Arjovsky et al. , 2019 ) . With these groups they train a new model with GDRO or IRM . Nam et al . ( 2020 ) ; Liu et l. ( 2021 ) observe that ERM models often misclassify data points in minority groups , and thus train another model with re-weighted or upsampled points misclassified by an initial ERM model . While these methods promisingly leverage ERM learned biases to significantly improve worst-group error without training group labels , there is still a gap between their robust performance and methods ’ such as GDRO that use group labels . In this work , we ask how else we can improve model robustness using a trained ERM model , and aim to close this gap by focusing on improving the learned representations of the robust model in the second stage . We support this direction with two key motivations . First , we find that higher worst-group performance consistently correlates with hidden-layer representations exhibiting higher dependence on class labels than spurious attributes . We quantify this correlation using geometric representation alignment ( Wang & Isola , 2020 ) , which measures the closeness of samples with the same class but different spurious attributes in the model feature space , and mutual information . This relation consistently holds across various data sets , and explains when prior upweighting methods improve worst-group error over ERM ( Fig . 4 ) . Second , we theoretically show that a model ’ s representation alignment for a given class can be used to upper bound the gap between its worst-group and average loss for that class . Thus , if we can improve representation alignment for a class , we can reduce the gap between worst-group and average loss for that class . We thus propose Correct-N-Contrast ( CNC ) , a two-stage procedure using contrastive learning to encourage better representation alignment within each class . In the first stage , we train a regularized ERM model similar to prior work ( Liu et al. , 2021 ; Creager et al. , 2021 ) , under the premise that ERM predictions help infer group information ( i.e. , spurious attributes ) . In the second stage , we wish to improve representation alignment by “ pulling together ” same-class datapoints and “ pushing apart ” different-class datapoints , regardless of their individual groups or spurious features . To do so via supervised contrastive learning , we use the heuristic that samples with the same ERM predictions exhibit similar spurious features ( and vice versa ) . With a randomly sampled anchor , we select samples with the same class but different ERM predictions as “ positives ” we want to pull together , and samples from different classes but the same ERM prediction as hard “ negatives ” we want to push apart . Training a second model with this sampling scheme and supervised contrastive learning encourages this model to ignore spurious correlations that the initial ERM model learned , and improves representation alignment between same-class data points . Thus , CNC corrects for the ERM model ’ s mistakes with contrastive learning in the second model . We evaluate CNC on four popular and diverse spurious correlation benchmarks . Among methods that similarly do not assume training group labels , CNC substantially improves worst-group accuracy , obtaining up to 7.7 % absolute lift ( from 81.1 % to 88.8 % on CelebA ) over the prior state-of-the-art JTT ( Liu et al. , 2021 ) , and averaging 3.4 % lift across the four tasks . We also find that CNC nearly closes the gap in worst-group accuracy with robust training methods that assume training group labels , only falling short of GDRO ’ s worst-group accuracy by 0.8 % absolute . Finally , we validate that CNC indeed reduces the alignment loss compared to prior methods . This corresponds to an up to 71.1 % smaller gap between worst-group versus average accuracy for data points in the same class . Contributions . We summarize our contributions as follows : 1 . We empirically show that a model ’ s worst-group performance correlates with the model ’ s alignment loss between different groups within a class , and analyze this connection theoretically . 2 . We propose CNC , a two-stage contrastive approach to improve representation alignment and thereby learn representations robust to spurious correlations . 3 . We validate that CNC significantly improves worst-group accuracy over existing methods on various benchmarks , and learns better-aligned representations less reliant on spurious features . 2 PRELIMINARIES . Problem setup . We present our setting and the loss objectives following Sagawa et al . ( 2019 ) . Let X = { x1 , . . . , xn } and Y = { y1 , . . . , yn } be a training data set of size n. Each data point has an observed feature vector xi ∈ X , label yi ∈ Y , and unobserved spurious attribute ai ∈ A . The set of groups G is defined as the set of all combinations of class label and spurious attribute pairs , i.e . G = Y ×A . Let C = |Y| be the number of classes and K = |G| be the number of groups . Following the classical supervised learning setting , we assume that each example ( xi , yi , ai ) is drawn from an unknown joint distribution P . We assume that at least one sample from each group is observed in the training data . Let Pg be the distribution conditioning on ( y , a ) = g , for any g ∈ G. Given a model fθ : X 7→ RC and a convex loss ` : X × Y 7→ R , let the worst-group loss be : Lwg ( fθ ) : = max g∈G E ( x , y , a ) ∼Pg [ ` ( fθ ( x ) , y ) ] . ( 1 ) ERM minimizes the training loss as a surrogate for the expected population loss Lavg : Lavg ( fθ ) : = E ( x , y , a ) ∼P [ ` ( fθ ( x ) , y ) ] ( 2 ) While ERM is the standard way to train neural nets , spurious correlations often cause ERM to obtain high error on minority groups even when average error is low . Group DRO , which minimizes the empirical version of ( 1 ) , is recognized as a strong baseline for improving worst-group error when the group labels { a1 , . . . , an } are available during training ( Sagawa et al. , 2019 ) . In contrast , we focus on the more challenging setting in which the group labels are not available during training . Contrastive learning . We briefly describe contrastive learning ( Chen et al. , 2020 ) , a central component of our approach . Let fθ be a neural network model with parameters θ . Let the encoder fenc : X 7→ Rd be the feature representation layers of fθ . Let fcls : Rd 7→ RC be the classification layer of fθ , which maps encoder representations to one-hot label vectors . We learn fenc with the supervised contrastive loss Lsupcon proposed in Khosla et al . ( 2020 ) . For each anchor x , we sample M positives { x+i } Mi=1 and N negatives { x − i } Ni=1 . Let y , { y + i } Mi=1 , { y − i } Ni=1 be the labels and z , { z+i } Mi=1 , { z − i } Ni=1 be the normalized outputs of fenc ( x ) for the anchor , positives , and negatives respectively . With input x mapped to z , the training objective for the encoder is to minimize : Lsupcon ( x ; fenc ) = E x , { x+i } Mi=1 , { x − i } Nj=1 [ − log exp ( z > z+i /τ ) ∑M m=1 exp ( z > z+m/τ ) + ∑N n=1 exp ( z > z−n /τ ) ] ( 3 ) where τ > 0 is a scalar temperature hyperparameter . Minimizing Eq . 3 leads to z being closer to z+ than z− in feature space . See Sec . 6 for further references related to contrastive learning . 3 MOTIVATIONS FOR REPRESENTATION ALIGNMENT . To motivate our method , we present our core observation that a model ’ s worst-group accuracy correlates with how well its learned representations depends on the class labels , but not the spurious attributes . First , we empirically observe that ERM learns spurious correlations by inspecting their hidden layer representations on several spuriously correlated data sets . We find that ERM ’ s worstgroup performance is inversely related to a cross-group alignment loss ( cf . Eq . ( 4 ) below ) and mutual information metrics . Second , we theoretically prove that this alignment loss serves as an upper bound on the gap between the average-group loss and the worst-group loss ( cf . Theorem 3.1 ) .

Summary: This paper focuses on improving the model robustness to group shifts without prior group information and bridges the gap to those methods with access to group labels (i.e., GDRO). It identifies the relation between the worst-group performance and representation alignment both empirical and theoretically, which motivates a contrastive approach for improving representation alignment and robustness. Empirically, the proposed method demonstrates improved worst-group performance over existing baselines.