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 considers the goal of training classifiers that achieve strong worst-case performance across groups with different spurious features, without assuming access to supervision based on the spurious features. The paper proposes a two-step method, which fist trains a model with ERM, then re-trains the model using a supervised contrastive approach where positive and negative samples are selected on the basis of misclassification characteristics of the ERM model.    The method addresses an important problem, and is well motivated & evaluated. In all this is a nice piece of work and I am cautiously happy to recommend its acceptance at this point. I do, however, have a couple of questions I would like to see some answers to (see below). Depending on the answers to these questions, and discussion with other reviewers, I am happy to consider raising my score.