Correct- N -Contrast: A Contrastive Approach for Improving Robustness to Spurious Correlations
Abstract
We propose Correct- N -Contrast (C N C), a contrastive learning method to improve robustness to spurious correlations when training group labels are unknown.
Our motivating observation is that worst-group performance is related to a representation alignment loss, which measures the distance in feature space between different groups within each class.
We prove that the gap between worst-group and average loss for each class is upper bounded by this alignment loss for that class.
Thus, C N C aims to improve representation alignment via contrastive learning.
First, C N C uses an ERM model to infer the group information.
Second, with a careful sampling scheme, C N C trains a contrastive model to encourage similar representations for groups in the same class.
We show that C N C signiﬁcantly improves worst-group accuracy over existing state-of-the-art methods on popular benchmarks, e.g., achieving 7 .
7% absolute lift in worst-group accuracy on the
CelebA dataset, and performs almost as well as methods trained with group labels.
C N C also learns better-aligned representations between different groups in each class, reducing the alignment loss substantially compared to prior methods.
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, a model may learn to classify animals solely based on background [35, 6, 19].
In this work, we focus on improving a model’s robustness to spurious correlations via its learned representations .
We support this direction with two key motivations.
First, we ﬁnd that higher worst-group performance consistently correlates with hidden-layer representations exhibiting higher dependence on class labels than spurious attributes.
We quantify this via a representation alignment loss [47], which measures the closeness of sample representations with the same class but different spurious attributes, and mutual information.
Second, we theoretically show that this alignment loss for a class can upper bound the worst-group vs. average loss gap for that class.
Thus, if we can reduces alignment loss for a class, we can reduce the gap between worst-group and average loss for that class.
We thus propose Correct- N -Contrast (C N C), a two-stage procedure using contrastive learning to encourage better representation alignment within each class.
First, we train an ERM model similar to prior work [25, 12], under the premise that ERM predictions help infer group information (i.e., spurious attributes).
Second, we improve representation alignment by “pulling together” same-class datapoints and “pushing apart” different-class datapoints, regardless of their spurious features.
To do so via contrastive learning, we reason 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” to pull together, and samples from different classes
(c)
[Figure]
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
Ours
[Figure] [Figure]
but the same ERM prediction as hard “negatives” to push apart. Training a second model with this sampling scheme and a contrastive loss encourages this model to ignore spurious correlations that the initial ERM model learned, and improves representation alignment between same-class data points.
Thus, C N C corrects for the ERM model’s mistakes with contrastive learning in the second model.
We evaluate C N C on four popular and diverse spurious correlation benchmarks.
Among methods that similarly do not assume training group labels, C N C substantially improves worst-group accuracy, obtaining up to 7.7% absolute lift over the prior state-of-the-art J TT [25] (from 81.1% to 88.8% on
CelebA), and averaging 3.4% lift across the four tasks.
We also ﬁnd that C N C nearly closes the gap in worst-group accuracy with methods that assume training group labels such as group DRO (GDRO) [37], only falling short of GDRO’s worst-group accuracy by 0.3% absolute.
Finally, we validate that
C N C 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.
Preliminaries
Problem setup.
Our setting follows that of Sagawa et al.
Let X = { x 1 , . . .
, x n } and Y = { y 1 , . .
, y n } be a training dataset of size n .
Each data point has an observed feature vector x i ∈ X , label y i ∈ Y , and unobserved spurious attribute a i ∈ A .
The set of groups G is deﬁned 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.
We assume that each example ( x i , y i , a i ) is drawn from an unknown joint distribution P , with at least one sample from each group observed in the training data.
Let P g be the distribution conditioning on ( y, a )
=
g , for any g ∈ G .
Spurious correlations can cause ERM to obtain high error on minority groups even when average error is low.
We tackle the setting where group labels are not available during training.
Contrastive learning.
We brieﬂy describe contrastive learning [9] as it applies here.
Let f θ be a neural network with parameters θ .
Let the encoder f enc : X (cid:55)→ R d be the feature representation layers of f θ .
Let f cls : R d (cid:55)→ R C be the classiﬁcation layer of f θ , which maps encoder representations to one-hot label vectors.
We learn f enc with a supervised contrastive loss L supcon [20].
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 f enc ( x ) for the anchor, positives, and negatives respectively.
With input x mapped to z , the training objective for the encoder is to minimize:
[Equation]
where τ > 0 is a scalar temperature hyperparameter.
We discuss further related work in Appendix C.
Motivations for representation alignment 3.1 Empirical relation between worst-group performance and representation alignment
Experimental setup .
We model spurious correlations with CMNIST ∗ , a colored MNIST dataset inspired by [3].
There are 5 digit classes and 5 colors.
We color a fraction p corr of the training samples with a color a associated with each class y , and color the test samples uniform-randomly.
To analyze learned representations, we train a LeNet-5 CNN [23] with ERM to predict digit classes, and inspect the outputs of the last hidden layer z
= f enc ( x ) using the following representation metrics.
Representation metrics .
First, we compute an alignment loss ˆ L align ( f enc ; g, g (cid:48) ) for two groups g =
( y, a ) and g (cid:48) =
( y, a (cid:48) )
where a (cid:54) = a (cid:48) .
This measures how closely f enc maps samples with the same class, but different spurious attributes, via Euclidean distance.
Letting G and G (cid:48) be the training data subsets in groups g and g (cid:48) respectively, and x and x (cid:48) be any two samples in G and G (cid:48) , we deﬁne:
[Equation]
Thus, lower ˆ L align means better alignment.
We also quantify representation dependence by estimating the mutual information (MI) of a model’s learned representations with the class label, i.e. ˆ I ( Y ; Z ) and the spurious attributes ˆ I ( A ; Z ) .
We defer computational details to Appendix D.
Results .
In Fig.
2 we show a strong association between worst-group error and alignment and mutual information.
As p corr increases, ERM models not only drop in worst-group accuracy, but also incur higher alignment loss (Fig. 2ab).
Fig.
2c further illustrates this with mutual information.
We plot the estimated mutual information and worst-group accuracy for models at each epoch.
High worst-group accuracy (>80%) corresponds to a combination of both high
ˆ I ( Y ; Z ) and low ˆ I ( A ; Z ) .
Fig.
2d also captures this trend with a trade off between high
ˆ I ( Y ; Z ) with ˆ I ( A ; Z ) as p corr increases (Fig. 2a).
[Figure]
3.2 Theoretical analysis of the alignment loss
The empirical observations in Fig.
2 suggest that lower alignment loss correlates with lower worstgroup error.
Next, we show that this connection applies much more generally.
We show that the maximum of ˆ L align ( f enc ; g, g (cid:48) ) , over any two groups g, g (cid:48) within the same class, can be used to upper bound the gap between the worst-group loss and average loss for that class.
We set up several notations before stating the result.
For any class label y ∈ Y , let G y be the set of groups with label y in G .
Deﬁne the worst-group loss among groups in G y as L wg ( f θ ; y ) :
= max g ∈G y E ( x,y,a ) ∼ P g [ (cid:96) ( f θ ( x ) , y )] and the average loss among groups in G y as L avg ( f θ ; y ) :
=
E ( x,y,a ) ∼ P : ∀ a ∈A [ (cid:96) ( f θ ( x ) , y )] .
Additionally, let L align ( f enc ; y ) be the class-speciﬁc alignment loss for groups in G y .
Finally let L align ( f enc ; y ) be the largest cross-group alignment loss among groups in G y :
[Equation]
Our main result is that ˆ L align ( f θ ; y ) is an upper bound on the gap between L wg ( f θ ; y ) and L avg ( f θ ; y )
(up to a norm multiplier and a concentration error), for any y ∈ Y .
Theorem 3.1 (Alignment loss upper bounds the gap between worst-group and average-group loss) .
In the setting described above, let f θ be any neural network satisfying that the weight matrix of the linear classiﬁcation layer W in f cls satisﬁes that (cid:107) W (cid:107) 2 ≤ B , for some constant B .
Let n g be the size of any group g ∈ G in the training data set.
Assume that the loss function (cid:96) ( x, y ) is 1 -Lipschitz in x and bounded from above by one.
Then, with probability at least 1 − δ over the randomness of the training data set samples, for any class y ∈ Y , the following holds:
[Equation]
The proof of Theorem 3.1 is deferred to Sec.
A. Since we also know that L avg ( f θ ; y )
≤ L wg ( f θ ; y ) , the result implies that in order to reduce the gap between the worst-group loss and the average loss for class y , it sufﬁces to reduce the alignment loss ˆ L align ( f θ ; y ) .
Correct- N -Contrast (C N C)
C N C involves two stages to improve robustness to spurious correlations.
In Stage ERM training , we train an ERM model f
ˆ θ on the training dataset { ( x i , y i ) } ni =1 and save predictions { ˆ y i } ni =1 .
While this model is trained to predict the ground-truth class labels, we reason that because it is susceptible to learning spurious correlations, its predictions will actually be good proxies for the spurious attributes of the datapoints. We then use these inferred spurious labels on the training data for the next stage (regularizing f ˆ θ to prevent memorization).
In Stage 2: Contrastive learning , we train a robust model with supervised contrastive learning, using { ˆ y i } ni =1 and two components:
Contrastive batch sampling.
We sample points such that by maximizing the similarity between anchors and positives (and keeping anchors and negatives apart), the Stage 2 model “ignores” spurious correlations.
With prediction set { ˆ y i } ni =1 , for each batch we randomly sample an anchor x i ∈ X (with label y i and ERM prediction ˆ y i ), M positives with the same class as y i but different ERM prediction than ˆ y i , and N negatives with different classes as y i but the same ERM prediction as ˆ y i .
Optimization objective and updating procedure.
We wish to learn aligned representations via contrastive learning while also classifying datapoints correctly.
As we have the training class labels, we jointly update both the model’s encoder layers f enc with a standard contrastive loss, and the full model f θ with a cross-entropy loss (with λ ∈ [0 , 1] as a balancing hyperparameter): [Equation] [Equation]
[Equation]
We summarize C N C in Algorithm 1, and include further implementation details in Appendix B. [Table]
Experimental results
We conduct experiments to answer the following questions: (1) Does C N C improve worst-group performance over prior state-of-the-art methods on spurious correlation benchmarks?
(2) Does C N C encourage learning network representations with greater alignment and class-label-only dependence?
5. C N C improves worst-group performance
We evaluate C N C on image classiﬁcation and NLP datasets with spurious correlations (Table 1).
We compare against ERM, GDRO, and comparable methods to tackle spurious correlations without training group labels: CVaR DRO [24], G EORGE [40], Learning from Failure (LfF) [30], Environment
Inference for Invariant Learning (EIIL) [12], Contrastive Input Morphing (CIM) [42], and Just Train
Twice (J TT ) [25] (Table 1).
C N C achieves highest worst-group accuracy among methods without training group labels on the CMNIST ∗ , Waterbirds and CelebA datasets, and near-SOTA worst-group accuracy on CivilComments.
While LfF, G EORGE , EIIL, and J TT similarly use a trained ERM model to estimate groups, C N C uniquely uses ERM predictions to encourage the robust model to learn desirable representations via contrastive learning. We reason that with this approach, by sampling positives and negatives from the ERM predictions, C N C more directly encourages the robust model to ignore learnable spurious correlations compared to previous invariant learning, input transformation, or upweighting approaches. We include additional evidence of this via GradCAM visualizations in Appendix F.6.
[Table]
C
N C learns representations less reliant on spurious features To shed light on C N C’s worst-group accuracy gains
[Table]
we investigate if models trained with C N C actually learn representations with higher alignment. Compared to ERM and J TT (the next-best performing method that does not require subgroup labels), C N C learns representations with signiﬁcantly higher alignment (lower alignment loss) and lower mutual information with spurious attributes (while having comparable mutual information with class labels)
(Fig. 3).
We ﬁnd that C N C representations exhibit the lowest alignment loss consistently for these data sets; this also corresponds to C N C models achieving the highest worst-group accuracy. Furthermore, while all methods result in representations that exhibit high mutual information with the class label (Fig. 3b), only C N C results in representations that drastically reduce mutual information with spurious attributes (Fig. 3c). In Fig. 4, we also illustrate this result on the Waterbirds data set via UMAP visualizations of the learned representations.
Notably, all training methods result in representations separable by class label. Yet ERM models exhibit strong separability by spurious attributes, and J TT models interestingly also still depict some learned dependency on the spurious attribute. However, C N C uniquely learns representations that strongly depict class-label-only dependence.
[Figure]
Conclusion
We present C N C, a two-stage contrastive approach that learn representations robust to spurious correlations.
We theoretically analyze the connection between alignment and worst-group vs. averagegroup losses. We show that C N C improves the robustness of learned representations by making them more class-dependent and less spurious-attribute-dependent, and that C N C achieves SOTA or near-SOTA worst-group performance across several benchmark datasets.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
[Yes] Our contributions and claims are outlined in the introduction, and reﬂected in Sections 3.1, 3.2, and 5.
(b) Did you describe the limitations of your work?
[Yes] See Appendix E. (c) Did you discuss any potential negative societal impacts of your work?
[Yes] See Appendix E. (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] See Sec 3.
(b) Did you include complete proofs of all theoretical results?
[Yes] See Appendix A 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] See link in Appendix G (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] See Appendix D (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] See Appendix F. (d) Did you include the total amount of compute and the type of resources used (e
g., type of GPUs, internal cluster, or cloud provider)?
[Yes] Hardware resource details in Appendix D. 4.
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[Yes] Licenses for code and datasets mentioned in the appendix where applicable.
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes] Code linked in Appendix G (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
A Omitted Proofs from Section 3.
In this section, we prove that within any class, the gap between the worst-group error and the average error can be upper bounded by the alignment loss times the Lipschitz constant, plus another concentration error term.
Proof of Theorem 3.1.
Consider two arbitrary groups, denoted by g 1 =
( y, a 1 ) and g 2 =
( y, a 2 ) , whose class labels are both y ∈ Y , whose spurious attributes are a 1 ∈ A and
a 2 ∈ A such that a 1 (cid:54) =
a 2 .
Let G 1 and G 2 be the subset of training data that belong to groups g 1 and g 2 , respectively.
We note that both G 1 and G 2 are non-empty since we have assumed that (in Section 2) there is at least one sample from each group in the training data set.
Let n g 1 =
| G 1 | and n g 2 =
| G 2 | be the size of these two groups, respectively.
Recall that f enc denotes the mapping of the encoder layers of the full neural network model f θ .
Since the classiﬁcation layer f cls is a linear layer, we have used W to denote the weight matrix of this layer.
Our deﬁnition of the cross-group alignment loss in equation (3), denoted as ˆ L align ( f θ ; y ) , implies that for g 1 and g 2 ,
[Equation]
Next, let E ( x,y,a 1 ) ∼P g
[ L avg ( Wf enc ( x ) , y )] be the average loss conditioning on a data point being sampled from group g 1 (and similarly for group g 2 ).
Let ∆( g 1 , g 2 ) be the difference between the population average losses:
[Equation]
Recall that G y ⊆
G is the set of groups that have class label y .
Since the loss (cid:96) ( · ) is bounded above by 1 according to our assumption, and is at least zero, by the Hoeffding’s inequality, the following result holds with probability at least 1 − δ , for all |G y | groups g ∈ G y ,
[Equation]
Thus, with probability at least 1 − δ , the following holds for any g 1 and g 2 in class y
(but having different spurious attributes)
[Equation]
Next
[Equation]
we focus on the RHS of equation (8).
First, equation (8) is also equal to the following:
[Equation]
Since we have also assumed that the loss function (cid:96) ( x
[Equation]
y ) is 1 -Lipschitz in x
[Equation]
the above is at most:
[Equation]
n g(cid:88) ( x,y,a 1 ) ∈ G(cid:88)
( x (cid:48) ,y,a 2 ) ∈ G(cid:107) Wf enc ( x ) − Wf enc ( x (cid:48) ) (cid:107)(since y is the same for x, x (cid:48) ) ≤ B n
[Equation]
(cid:88) ( x,y,a 1 )
[Equation]
G(cid:88) ( x (cid:48) ,y,a 2 )
[Equation]
G(cid:107) f enc ( x ) [Equation] f enc ( x (cid:48)
[Equation]
(cid:107)(because (cid:107) W (cid:107) 2 ≤ B as assumed) ≤ B · ˆ L align ( f θ y ) .
(because of equation (6))
Thus, we have shown that for any g 1 and g 2 within class y ,
[Equation]
Finally, we use the above result to bound the gap between the worst-group loss and the average loss.
For every group g ∈ G , let p g denote the prior probability of observing a sample from P in this group.
Let q y = (cid:80) g (cid:48)
∈G y p g (cid:48) .
Let h ( g ) be a short hand notation for
[Equation]
The average loss among the groups with class label y is L avg ( f θ ; y )
= (cid:80)
g ∈G y p g q y h ( g ) .
The worstgroup loss among the groups with class label y is L wg ( f θ ; y )
=
max g ∈G y h ( g ) .
Let g (cid:63) be a group that incurs the highest loss among groups in G y .
We have L wg ( f θ ; y ) − L avg ( f θ ; y ) is equal to
[Equation]
The last step uses equation (9) on ∆( g (cid:63) , g ) and the fact that q y = (cid:80) g (cid:48) ∈G y p g (cid:48) .
Thus, we have shown that the gap between the worst-group loss and the average loss among the groups with the same class label is bounded by the above equation.
The proof is now complete.
As a remark, the proof of Theorem 3.1 can be adapted to the case when the loss function (cid:96) ( · ) is
Lipschitz for some ﬁxed constant and when the loss function is bounded from above by some ﬁxed constant as well.
We leave the details to the reader.
The astute reader will note that Theorem 3.1 focuses on comparing groups within the same class y , for any y ∈ Y .
A natural follow-up question is what happens when comparing across groups with different labels.
Let L
wg ( f θ ) = max y ∈Y L
wg ( f θ ; y ) be the worst-group loss across all the labels.
Recall that L avg ( f θ ) is the average loss for the entire population of data.
We generalize Theorem 3. to this setting in the following result.
Corollary A.1 (Extension of Theorem 3.1 to compare across different classes) .
In the setting of
Theorem 3.1, let q y = (cid:80) g ∈G
y p g be the prior probability of observing a sample drawn from P with label y , for any y ∈ Y .
We have that with probability at least 1 − δ , the following holds:
[Equation]
Proof.
We generalize the argument in the previous result to compare across different labels.
The worst-group loss across different labels is
[Equation]
of equation
[Equation]
[Equation] [Equation]
Since (cid:80) g ∈G p
g h ( g ) =
L avg ( f θ ) , we thus conclude that [Equation] n
[Equation]
The proof is now complete.
An example showing that Corollary A.1 is tight.
We describe a simple example in which the factor (cid:16) min y ∈Y q y (cid:17) − 1 in equation (13) is tight (asymptotically).
Suppose there are k perfectly balanced classes so that q y = 1 /k , for every y ∈ Y .
There is one data point from each class, with loss equal to 0 for all except one of them.
The worst-group loss is 1 whereas the average loss is 1 /k .
Thus, there is a factor of k between the worst-group loss and the average loss.
For equation (13), the factor
[Equation]
since q y = 1 /k for every y ∈ Y in this example.
Thus, this factor matches the (multiplicative) factor between the worst-group loss and the average loss in this example.
B Contrastive algorithm design details
In this section, we provide further details on the training setup and contrastive batch sampling, pseudocode, and additional properties related to C N C’s implementation.
B.1 Algorithm
Correct
Let ˆ y be the predicted (
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
B.2 Training setup
In Fig.
5, we illustrate the two training stages of Correct- N -Contrast described in Sec. 4.
In Stage 1, we ﬁrst train an ERM model with a cross-entropy loss.
For consistency with Stage 2, we depict the output as a composition of the encoder and linear classiﬁer layers.
Then in Stage 2, we train a new model with the same architecture using contrastive batches sampled with the Stage 1 ERM model and a supervised contrastive loss (1) (which we compute after the depicted representations are ﬁrst normalized) to update the encoder layers.
Note that unlike prior work in contrastive learning [9, 20], as we have the class labels of the anchors, positives, and negatives, we also continue forward-passing the unnormalized representations (encoder layer outputs) and compute a cross-entropy loss to update the classiﬁer layers while jointly training the encoder layers as well.
We also note that unlike prior work, we wish to learn invariances between anchors and positives that maximally reduce the presence of features not needed for classiﬁcation.
We thus do not pass the representations through an additional projection network [9].
Instead, we use Eq.
1 to compute the supervised contrastive loss directly on the encoder outputs z = f enc ( x ) .
In Appendix ??
, we study ablations with both design choices.
“Waterbird” 2048-D 2-D
Cross-entropy
Stage 1: ERM training “Waterbird”
Cross-entropy
Stage 2: Correct- N -Contrast “Waterbird” “Landbird”
Contrastive Contrastive
Input
Positive
Anchor
Negative
Loss Function
Forward Pass
Backward Pass
B.3 Two-sided contrastive batch implementation
We provide more implementation details on our default contrastive batch sampling approach described in Sec. 4.
For additional contrastive signal per batch, we can double the pairwise comparisons in a training batch by switching the anchor and positive roles.
This is similar to the NT-Xent loss in prior contrastive learning work [9].
We switch the role of the anchor and ﬁrst positive sampled in a contrastive batch, and sample additional positives and negatives using the same guidelines but adjusting for the “new” anchor.
We denote this as “two-sided” sampling in contrast with the “one-sided” comparisons we get with just the original anchor, positives, and negatives.
Implementing this sampling procedure in practice is simple.
First, recall our initial setup with trained
ERM model f ˆ θ , its predictions { ˆ y i } ni =1 on training data { ( x i , y i ) } ni
=1 (where ˆ y i
=
f ˆ θ ( x i ) ), and number of positives and negatives to sample M and N .
We then sample batches with Algorithm 3.
Require
Because the initial anchors are then datapoints that the ERM model gets correct, under our heuristic we infer { x i } Mi =1 as samples from the majority group.
Similarly the M positives { x + m } Mm =1 and N negatives { x − n } Nn =1 that it gets incorrect are inferred to belong to minority groups.
For one batch, we then compute the full contrastive loss with
[Equation]
supcon
[Equation]
f
[Equation]
(14) where e.g. ˆ L supcon (cid:0) x 1 , { x + m } Mm =1 , { x − n } Nn =1 ; f enc (cid:1) is given by:
[Equation]
and again let z be the normalized output f enc ( x ) for corresponding x .
We compute the cross-entropy component of the full loss for each x in the two-sided batch with its corresponding label y .
[Figure]
B.4 Summary of C N C design choices and properties
We summarize C N C’s design choices, additional properties, and differences from standard supervised contrastive learning below.
In Appendix F.4, we empirically validate each component.
No projection network.
As we wish to learn data representations that maximize the alignment between anchor and positive datapoints, we do not compute the contrastive loss with the outputs of an additional nonlinear projection network.
This is inspired by the logic justifying a projection head in prior contrastive learning, e.g. SimCLR [9], where the head is included because the contrastive loss trains representations to be “invariant to data transformation” and may encourage removing information “such as the color or orientation of objects”.
In our case, we view inferred datapoints with the same class but different spurious attributes as “transformations” of each other, and we hypothesize that maximally removing these differences can help us improve worst-group performance.
Two-sided contrastive sampling.
To incorporate additional comparisons between datapoints that only differ in spurious attribute during training, we employ “two-sided” contrastive batch sampling.
This lets us equally incorporate instances where the second contrastive model in C N C treats datapoints that the initial ERM model got incorrect and correct as anchors.
Additional intrinsic hard positive/negative mining.
Because the new model corrects for potentially learned spurious correlations by only comparing and contrasting datapoints that differ in class label or spurious attribute, but not both (as dictated by the initial ERM model’s outputs), the contrastive batches naturally carry “hard” positives and negatives.
Thus, our approach provides a natural form of hard negative mining (in addition to the intrinsic hard positive / negative mining at the gradient level with InfoNCE-style contrastive losses [9, 20]) while avoiding class collisions, two nontrivial challenges in standard self-supervised contrastive learning [36, 50, 10].
C Further related work discussion
We provide additional discussion of related work and connections to our work below.
C.1 Improving robustness to spurious correlations
Our core objective is to improve model robustness to group or subpopulation distribution shifts that arise from the presence of spurious correlations, speciﬁcally for classiﬁcation tasks.
Because these learnable correlations hold for some but not all samples in a dataset, standard training with
ERM may result in highly variable performance: a model that learns to classify datapoints based on these spurious correlations does well for some subsets or “groups” of the data but not others.
To improve model robustness and avoid learning spurious correlations, prior work introduces the goal to maximize the worst-group accuracy [37].
Related works broadly fall under two categories:
Improving robustness with group information.
If information such as spurious attribute labels is provided, one can divide the data into explicit groups as deﬁned in Sec. 2, and then train to directly minimize the worst group-level error among these groups.
This is done in group DRO (GDRO) [37], where the authors propose an online training algorithm that focuses training updates over datapoints from higher-loss groups.
Goel et al.
[15] also adopt this approach with their method CycleGAN
Augmented Model Patching (CAMEL).
However, similar to our motivation, they argue that a stronger modeling goal should be placed on preventing a model from learning group-speciﬁc features.
Their approach involves ﬁrst training a CycleCAN [53] to learn the data transformations from datapoints in one group to another that share the same class label.
They then apply these transformations as data augmentations to different samples, intuitively generating new versions of the original samples that take on group-speciﬁc features.
Finally they train a new model with a consistency regularization objective to learn invariant features between transformed samples and their sources.
Unlike their consistency loss, we accomplish a similar objective to learn group-invariant features with contrastive learning.
Our ﬁrst training stage is also less expensive.
Instead of training a CycleGAN and then using it to augment datapoints, we train a relatively simple standard ERM classiﬁcation model, sometimes with only a few number of epochs, and use its predictions to identify pairs of datapoints to serve a similar purpose.
Finally, unlike both CAMEL and GDRO, we do not require spurious attribute or group labels for each training datapoints.
We can then apply C N C in less restrictive settings where such information is not known.
Related to GDRO are methods that aim to optimize a "Pareto-fair" objective, more general than simply the worst-case group performance.
Notable examples are the works of Balashankar et al. [5] and Martinez et al.
However, these approaches similarly do not directly optimize for good representation alignment (unlike our work).
Improving robustness without training group information.
More similar to our approach are methods that do not assume group information at training time, and only require validation set spurious attribute labels for ﬁne-tuning.
As validation sets are typically much smaller in size than training sets, an advantage of C N C and comparable methods is that we can improve the accessibility of robust training methods to a wider set of problems.
One popular line of work is distributionally robust optimization (DRO), which trains models to minimize the worst loss within a ball centered around the observed distribution [7, 49, 14, 24, 13, 32].
This includes the CVaR DRO [24] method we evaluate against.
However, prior work has shown that these approaches may be too pessimistic, optimizing not just for worst-group accuracy but worst possible accuracy within the distribution balls [37], or too undirected, optimizing for too many subpopulations, e.g. by ﬁrst upweighting minority points but then upweighting majority points in later stages of training [25].
Pezeshki et al.
[34] instead suggest that gradient starvation (GS), where neural networks only learn to capture statistically dominant features in the data [11], is the main culprit behind learning spurious correlations, and introduce a “spectral decoupling” regularizer to alleviate GS.
However this does not directly prevent models from learning dependencies on spurious attributes.
Similar to CAMEL, [42] propose Contrastive Input
Morphing (CIM), an image dataset-speciﬁc method that aims to learn input feature transformations that remove the effects of spurious or task-irrelevant attributes.
They do so without group labels, training a transformation network with a triplet loss to transform input images such that a given transformed image’s structural similarity metric (based on luminance, contrast, and structure [48]) is more similar to a “positive” image from the same class than a “negative” image from a different class.
They then train a classiﬁer on top of these representations.
Instead of pixel-level similarity metrics, C N C enforces similarity in a neural network’s hidden-layer representations, allowing C N C to apply to non-image modalities.
Additionally, we sample positives and negatives not just based on class label, but also the learned spurious correlations of an ERM model (via its trained predictions).
We hypothesize that our sampling scheme, which intuitively provides "harder" positive and negative examples, allows C N C to more strongly overcome spurious correlations.
Most similar to our approach are methods that ﬁrst train an initial ERM model with the class labels as a way to identify data points belonging to minority groups, and subsequently train an additional model with greater emphasis on the estimated minority groups.
Sohoni et al.
[40] demonstrate that even when only trained on the class labels, neural networks learn feature representations that can be clustered into groups of data exhibiting different spurious attributes.
They use the resulting cluster labels as estimated group labels before running GDRO on these estimated groups.
Meanwhile, Nam et al.
[30] train a pair of models, where one model minimizes a generalized cross-entropy loss [51], such that the datapoints this model classiﬁes incorrectly largely correspond to those in the minority group.
They then train the other model on the same data but upweight the minority-group-estimated points.
While they interweave training of the biased and robust model, Liu et al.
[25] instead train one model ﬁrst with a shortened training time (but the standard cross-entropy objective), and show that then upsampling the incorrect data points and training another model with ERM can yield higher worst-group accuracy.
[12] ﬁrst train an ERM model, and then softly assign the training data into groups under which the initial trained ERM model would maximally violate the invariant risk minimization (IRM) objective.
In particular, the IRM objective is maximally satisﬁed if a model’s optimal classiﬁer is the same across groups [3], and EIIL groups are inferred such that the initial
ERM model’s representations exhibit maximum variance within each group.
Finally, Nagarajan et al.
[29] provides a theoretical understanding of how ERM picks up spurious features under dataset imbalance.
They consider a setting involve a single spurious feature that is correlated with the class label and analyze the max-margin classiﬁer in the presence of this spurious feature.
In our work, we demonstrate that the ERM model’s predictions can be leveraged to not only estimate groups and train a new model with supervised learning but with different weightings.
Instead, we can speciﬁcally identify pairs of points that a contrastive model can then learn invariant features between.
Our core contribution comes from rethinking the objective with a contrastive loss that more directly reduces the model’s ability to learning spurious correlations.
C.2 Contrastive learning
Our method also uses contrastive learning, a simple yet powerful framework for both self-supervised [9, 31, 43, 41, 39, 18, 36] and supervised [20, 16] representation learning.
The core idea is to learn data representations that maximize the similarity between a given input “anchor” and distinct different views of the same input (“positives”).
Frequently this also involves contrasting positives with “negative” data samples without any assumed relation to the anchor [4].
Core components then include some way to source multiple views, e.g. with data transformations [9], and training objectives similar to noise contrastive estimation [17, 28].
An important component of contrastive learning is the method by which appropriate positives and negatives are gathered.
For sampling positives, Chen et al.
[9] show that certain data augmentations (e.g. crops and cutouts) may be more beneﬁcial than others (e.g. Gaussian noise and Sobel ﬁltering) when generating anchors and positives for unsupervised contrastive learning.
[45] theoretically study how data augmentations help contrastive models learn core content attributes which are invariant to different observed “style changes”.
They propose a latent variable model for self-supervised learning.
Tian et al.
[44] further study what makes good views for contrastive learning.
They propose an “InfoMin principle”, where anchors and positives should share the least information necessary for the contrastive model to do well on the downstream task.
For sampling negatives, Robinson et al. [36] show that contrastive learning also beneﬁts from using “hard” negatives, which (1) are actually a different class from the anchor (which they approximate in the unsupervised setting) and (2) embed closest to the anchor under the encoder’s current data representation.
Both of these approaches capture the principle that if positives are always too similar to the anchor and negatives are always too different, then contrastive learning may be inefﬁcient at learning generalizable representations of the underlying classes.
In our work, we incorporate this principle by sampling data points with the same class label but different ERM predictions–presumably because of spurious attribute differences–as anchor and positive views, while sampling negatives from data points with different class labels but the same
ERM prediction as the anchor.
The anchors and positives are different enough that a trained ERM model predicted them differently, while the anchors and negatives are similar enough
that the trained
ERM model predicted them the same.
Contrasting the above then allows us to exploit both “hard” positive and negative criteria for our downstream classiﬁcation task.
In Appendix F.4.1, we show that removing this ERM-guided sampling and only sampling positives and negatives based on class information leads to substantially lower worst-group accuracy with C N C.
One limitation of our current theoretical analysis regarding the alignment loss (cf.
Section 3.2) is that we require knowing the group labels to compute the RHS of equation (4) (in particular, the alignment loss).
An interesting question for future work is to provide a better theoretical understanding of the alignment induced by C N C in the context of spurious correlations.
C.3 Learning invariant representations
Finally, our work is also similar in motivation to Invariant Risk Minimization (IRM) [3] and other related works in domain-invariant learning [22, 33, 1, 12].
These methods aim to train models that learn a single invariant representation that is consistently optimal (e.g. with respect to classifying data) across different domains or environments.
These environments can be thought of as data groups, and while traditionally methods such as IRM require that environment labels are known, recent approaches such as Environment Inference for Invariant Learning (EIIL) [12] similarly aim to infer environments with an initial ERM model.
In EIIL, they next train a more robust model with an invariant learning objective, similarly selecting models based on the worst-group error on the validation set.
However, their main goal for learning these invariances is to extrapolate to out-of-domain distribution shifts not seen during training, as opposed to improving worst-group performance / robustness to group shifts, which is made challenging by rare groups as in our setting.
D Additional experimental details
We ﬁrst further describe our evaluation benchmarks in Appendix D.1.
We next provide further details on how we calculate the reported metrics and the experimental hyperparameters of our main results in Appendix D.1.
For all methods, following prior work [25, 40, 30, 37, 12] we report the test set worst-group and average accuracies from models selected through hyperparameter tuning for the best validation set worst-group accuracy.
While different methods have different numbers of tunable hyperparameters, we try to keep the number of validation queries as close as possible while tuning for fair comparison.
D.1 Dataset details
Colored MNIST (CMNIST ∗ ).
We evaluate with a version of the Colored MNIST dataset proposed in Arjovsky et al.
The goal is to classify MNIST digits belonging to one of 5 classes
Y = { (0, 1), (2, 3), (4, 5), (6, 7), (8, 9) } , and treat color as the spurious attribute.
In the training data, we color p corr of each class’s datapoints with an associated color a , and color the rest randomly.
If p corr is high, trained ERM models fail to classify digits that are not the associated color.
We pick a from uniformly interspersed intervals of the hsv colormap, e.g. 0 and 1 digits may be spurious correlated with the color red ( #ff0000 ), while 8 and 9 digits may be spuriously correlated with purple ( #ff0018 ).
The full set of colors in class order are A = { #ff0000 , #85ff00 , #00fff3 , #6e00ff , #ff0018 } .
For validation and test data, we color each datapoint randomly with a color a ∈ A .
We use the default test set from MNIST, and allocate 80%-20% of the default MNIST training set to the training and validation sets.
For main results, we set p corr = 0 .
Waterbirds.
We evaluate with the Waterbirds dataset, which was introduced as a standard spurious correlations benchmark in Sagawa et al.
In this dataset, masked out images of birds from the
CUB dataset [46] are pasted on backgrounds from the Places dataset [52].
Bird images are labeled either as waterbirds or landbirds; background either depicts water or land.
From CUB, waterbirds consist of seabirds (ablatross, auklet, cormorant, frigatebird, fulmar, gull, jaeger, kittiwake, pelican, pufﬁn, tern) and waterfowl (gadwell, grebe, mallard, merganser, guillemot, Paciﬁc loon).
All other birds are landbirds.
From Places, water backgrounds consist of ocean and natural lake classes, while land backgrounds consist of bamboo forest and broadleaf forest classes.
The goal is to classify the foreground bird as Y =
{ waterbird, landbird } , where there is spurious background attribute A = { water background, land background } .
We use the default training, validation, and test splits [37], where in the training data 95% of waterbirds appear with water backgrounds and 95% of landbirds appear with land backgrounds.
Trained ERM models then have trouble classifying waterbirds with land backgrounds and landbirds with water backgrounds.
For validation and test sets, water and land backgrounds are evenly split among landbirds and waterbirds.
CelebA. We evaluate with the CelebA spurious correlations benchmark introduced in Sagawa et al.
The goal is to classify celebrities’ hair color Y = { blond, not blond } , which is spuriously correlated with the celebrity’s identiﬁed gender
A = { male, female } .
We use the same training, validation, test splits as in Sagawa et al.
Only 6% of blond celebrities are male; trained ERM models perform poorly on this group.
CivilComments-WILDS.
We evaluate with the CivilComments-WILDS dataset from Koh et al.
[21], derived from the Jigsaw dataset from Borkan et al.
Each datapoint is a real online comment curated from the Civil Comments platform, a commenting plugin for independent news sites.
For classes, each comment is labeled as either toxic or not toxic.
For spurious attributes, each comment is also labeled with the demographic identities { male, female, LGBTQ, Christian, Muslim, other religions, Black, White } mentioned; multiple identities may be mentioned per comment.
The goal is to classify the comment Y = { toxic, not toxic } .
As in Koh et al.
[21], we evaluate with
A = { male, female, LGBTQ, Christian, Muslim, other religions, Black, White } .
There are thentotal groups corresponding to (toxic, identity) and (not toxic, identity) for each identity.
Groups may overlap; a datapoint falls in a group if it mentions the identity.
We use the default data splits [21].
In Table 2, we list the percentage of toxic comments for each identity based on the groups.
Trained
ERM models in particular perform less well on the rarer toxic groups.
[Table]
2 Implementation details D.2.1 Reported metrics
Main results.
For the CMNIST ∗ , Waterbirds, and CelebA datasets, we run C N C with three different seeds, and report the average worst-group accuracy over these three trials in Table 1.
As we use the same baselines and comparable methods as Liu et al.
[25], we referenced their main results for the reported numbers, which did not have standard deviations or error bars reported.
For CivilCommentsWILDS, due to time and compute constraints we only reported one run.
We note that CMNIST ∗ here is extremely challenging, as minority groups together only make up 0 .
5% of the training set.
This severe imbalance explains the very poor worst-group performance of ERM (as well as a couple other methods that fail to sufﬁciently remediate issue).
Estimated mutual information.
We give further details for calculating the representation metric introduced in Sec. 3.
As a reminder, we report both alignment and estimated mutual information metrics to quantify how dependent a model’s learned representations are on the class labels versus the spurious attributes, and compute both metrics on the representations
Z = { f enc ( x ) } over all test set data points x .
Then to supplement the alignment loss calculation in Sec. 3, we also estimate I ( Y ; Z ) and I ( A ; Z ) , the mutual information between the model’s data representations and the class labels and spurious attribute labels respectively.
To ﬁrst estimate mutual information with Y , we ﬁrst approximate p ( y | z ) by ﬁtting a multinomial logistic regression model over all representations Z to classify y .
With the empirical class label distribution p ( y ) , we compute:
[Equation]
We do the same but substitute the spurious attributes a for y to compute ˆ I ( A ; Z ) .
D.2.2 Stage 1 ERM training details
We describe the model selection criterion, architecture, and training hyperparameters for the initial
ERM model in our method.
To select this model, recall that we ﬁrst train an ERM model to predict the class labels, as the model may also learn dependencies on the spurious attributes.
Because we then use the model’s predictions on the training data to infer samples with different spurious attribute values but the same class label, we prefer an initial ERM model that better learns this spurious dependency, and importantly also does not overﬁt to the training data.
Inspired by the results in prior work [40, 25], we then explored using either a standard ERM model, one with high (cid:96) -2 regularization ( weight decay = 1 ), or one only trained on a few number of epochs.
To select among these, because the validation data has both class labels and spurious attributes, we choose the model with the largest gap between worst-group and average accuracy on the validation set.
For fair comparison to JTT, we use the same batch size, learning rate, momentum, optimizer, default weight decay, and number of epochs as reported in [25] to obtain these models.
We detail the ERM architecture and hyperparameters for each dataset below:
Colored MNIST.
We use the LeNet-5 CNN architecture in the pytorch image classiﬁcation tutorial.
We train with SGD, few epochs E = 5 , SGD, learning rate 1e-3, batch size 32, default weight decay 5e- and momentum 0.9.
Waterbirds.
We use the torchvision implementation of ResNet-50 with pretrained weights from
ImageNet as in Sagawa et al.
Also as in [37], we train with SGD, default epochs E = 300 , learning rate 1e-3, batch size 128, and momentum 0.9.
However we use high weight decay 1 . 0 .
CelebA.
We also use the torchvision ImageNet-pretrained ResNet-50 and default hyperparameters from Sagawa et al.
[37] but with high weight decay: we train with SGD, default epochs E = 50 , learning rate 1e-4, batch size 128, momentum 0 .
9 , and high weight decay 0 .
CivilComments-WILDS.
We use the HuggingFace ( pytorch-transformers ) implementation of
BERT with pretrained weights and number of tokens capped at 300 as in Koh et al.
As in Liu et al.
[25], with other hyperparameters set to their defaults [21] we tune between using the AdamW optimizer with learning rate 1e-5 and SGD with learning rate 1e-5, momentum 0.9, and the PyTorch
ReduceLROnPlateau learning rate scheduler.
Based on our criterion, we use SGD, few number of epochs E = 2 , learning rate 1e-5, batch size 16, default weight decay 1e-2, and momentum 0 .
D.2.3 Contrastive batch sampling details
We provide further details related to collecting predictions from the trained ERM models, and the number of positives and negatives that determine the contrastive batch size.
ERM model prediction.
To collect trained ERM model predictions on the training data, we explored two approaches: (1) using the actual predictions, i.e. the argmax for each classiﬁer layer output vector, and (2) clustering the representations, or the last hidden-layer outputs, and assigning a cluster-speciﬁc label to each data point in one cluster.
This latter approach is inspired by Sohoni et al.
[40], and we similarly note that ERM models trained to predict class labels in spuriously correlated data may learn data representations that are clusterable by spurious attribute.
As a viable alternative to collecting the “actual” predictions of the trained ERM model on the training data, with C classes, we can then cluster these representations into C clusters, assign the same class label only to each data point in the same cluster, and choose the label-cluster assignment that leads to the highest accuracy on the training data.
We also follow their procedure to ﬁrst apply UMAP dimensionality reduction to 2 UMAP components, before clustering with K-means or GMM [40].
To choose between all approaches, we selected the procedure that lead to highest worst-group accuracy on the validation data after the second-stage of training.
While this cluster-based prediction approach was chosen as a computationally efﬁcient heuristic, we found that in practice it either lead to comparable or better ﬁnal worst-group accuracy on the validation set.
To better understand this, as a preliminary result we found that when visualizing the validation set predictions with the Waterbirds dataset, the cluster-based predictions captured the actual spurious attributes better than the classiﬁer layer predictions (Fig. 7).
We defer additional discussion to Sohoni et al.
[40] and leave further analysis to future work.
Number of positives and negatives per batch.
One additional difference between our work and prior contrastive learning methods [9, 20] is that we speciﬁcally construct our contrastive batches by sampling anchors, positives, and negatives ﬁrst.
This is different from the standard procedure of randomly dividing the training data into batches ﬁrst, and then assigning the anchor, positive, and
[Figure]
negative roles to each datapoint in a given batch.
As a result, we introduce the number of positives
M and the number of negatives N as two hyperparameters that primarily inﬂuence the size of each contrastive batch (with number of additional anchors and negatives also following M and N with two-sided batches).
To maximize the number of positive and negative comparisons, as a default we set M and N to be the maximum number of positives and negatives that ﬁt the sampling criteria speciﬁed under Algorithm 3 that also can ﬁt in memory.
In Appendix D.2.4, for each dataset we detail the ERM prediction method and number of positives and negatives sampled in each batch.
D.2.4 Stage 2 contrastive model training details
In this section we describe the model architectures and training hyperparameters used for training the second model of our procedure, corresponding the reported worst-group and average test set results in Table 1.
In this second stage, we train a new model with the same architecture as the initial ERM model, but now with a contrastive loss and batches sampled based on the initial ERM model’s predictions.
We report test set worst-group and average accuracies from models selected with hyperparameter tuning and early stopping based on the highest validation set worst-group accuracy.
For all datasets, we sample contrastive batches using the clustering-based predictions of the initial
ERM model.
Each batch size speciﬁed here is also a direct function of the number of positives and negatives: 2 M + 2 N .
Colored MNIST.
We train a LeNet-5 CNN.
For C N C, we use M = 32 , N = 32 , batch size 128 , temperature τ = 0 .
05 , contrastive weight λ = 0 .
75 , SGD optimizer, learning rate 1e-3, momentum 0.9, and weight decay 1e-4.
We train for 3 epochs, and use gradient accumulation to update model parameters every 32 batches.
Waterbirds.
We train a ResNet-50 CNN with pretrained ImageNet weights.
For C N C, we use
M = 17 , N = 17 , batch size 68 , temperature τ = 0 .
1 , contrastive weight λ = 0 .
75 , SGD optimizer, learning rate 1e-4, momentum 0.9, weight decay 1e-3.
We train for 5 epochs, and use gradient accumulation to update model parameters every 32 batches.
CelebA. We train a ResNet-50 CNN with pretrained ImageNet weights.
For C N C, we use M = 64 ,
N = 64 , batch size 256 , temperature τ = 0 .
05 , contrastive weight λ = 0 .
75 , SGD optimizer, learning rate 1e-5, momentum 0.9, and weight decay 1e-1.
We train for 15 epochs, and use gradient accumulation to update model parameters every 32 batches.
CivilComments-WILDS.
We train a BERT model with pretrained weights and max number of tokens 300 .
For C N C, we use M = 16 , N = 16 , batch size 64 , temperature τ = 0 .
1 , contrastive weight λ = 0 .
75 , AdamW optimizer, learning rate 1e-4, weight decay 1e-2, and clipped gradient norms.
We train for 10 epochs, and use gradient accumulation to update weights every 128 batches.
D.2.5 Comparison method training details
As reported in the main results (Table 1) we compare C N C with the ERM and Group DRO baselines, as well as robust training methods that do not require spurious attribute labels for the training data:
CVaR DRO [24], G EORGE [24], Learning from Failure (LfF) [24], Contrastive Input Morphing (CIM) [42], Environment Inference for Invariant Learning (EIIL) [12], and Just Train Twice (J TT ) [25].
For each dataset, we use the same model architecture for all methods.
For the Waterbirds, CelebA, and
CivilComments-WILDS datasets, we report the worst-group and average accuracies reported in Liu et al.
[25] for ERM, CVaR DRO, LfF, and J TT .
For G EORGE , we report the accuracies reported in Sohoni et al.
For CIM, we report results from Waterbirds and CelebA from [42] using the
CIM + variational information bottleneck implementation [2], which achieves the best worst-group performance in their results.
For EIIL, we report results from Waterbirds and CivilComments-WILDS from [12].
For these hyperparameters, we defer to the original papers.
For GDRO, we reproduce the results with the same optimal hyperparameters over three seeds.
For Colored MNIST, we run implementations for G EORGE , CIM, EIIL, J TT , and GDRO, using code from their authors respectively.
LfF and CVaR DRO are also run with code from the J TT authors.
We include training details for our own implementations below:
Colored MNIST CMNIST ∗ We run
all methods for epochs, reporting test set accuracies with early stopping.
For J TT , we train with SGD, learning rate 1e-3, momentum 0.9, weight decay 5e-4, batch size 32.
We use the same initial ERM model as C N C, with hyperparameters described in Appendix D.2.2.
For upsampling we ﬁrst tried constant factors { 10 , 100 , 1000 } .
We also tried a resampling strategy where for all the datapoints with the same initial ERM model prediction, we upsample the incorrect points such that they equal the correct points in frequency, and found this worked the best.
With p corr = 0 .
995 , this upsamples each incorrect point by roughly 1100 .
We also use this approach for the results in Fig. 9.
For GDRO we use the same training hyperparameters as J TT , but without the upsampling and instead set group adjustment parameter C = 0 .
For LfF, we use the same hyperparameters as J TT , but instead of upsampling gridsearched the q parameter ∈ { 0 . 1 , 0 . 3 , 0 .
7 , 0 . 9 } , using q = 0 .
For CVaR DRO we do the same but use hyperparameter α = 0 .
For G EORGE we train with SGD, learning rate 1e-3, momentum 0.9, weight decay 5e-4.
For CIM, we use the CIM + VIB implementation.
We train with SGD, learning rate 1e-3, weight decay 5e-4, β parameter 10 , and λ parameter 1e-5.
For EIIL, for environment inference we use the same initial ERM model as C N C and J TT , and update the soft environment assignment distribution with Adam optimizer, learning rate 1e-3, and 10000 steps.
Following [12]’s own colored MNIST experiment, we train the second model with IRM, using learning rate 1e-2, weight decay 1e-3, penalty weight 100 , and penalty annealing parameter 80 .
CelebA. We also tune EIIL for CelebA. We again use the same initial ERM model as C N C, and update the soft environment assignment distribution with Adam optimizer, learning rate 1e-3, and 10000 steps.
We train the second model with GDRO, using SGD, 50 epochs, learning rate 1e-5, batch size 128, weight decay 0 . 1 , and group adjustment parameter 3 .
In Section D.2.6, we include our sweeps for both the method-speciﬁc and general hyperparameters in
CIM, EIIL, J TT , and C N C.
Comparison limitations.
One limitation of our comparison is that because for each dataset we sample new contrastive batches which could repeat certain datapoints, the number of total batches per epoch changes.
For example, 50 epochs training the second model in C N C does not necessarily lead to the same total number of training batches as 50 epochs training with ERM, even if they use the same batch size.
However, we note that the numbers we compare against from Liu et al.
[25] are reported with early stopping.
In this sense we are comparing the best possible worst-group accuracies obtained by the methods, not the highest worst-group accuracy achieved within a limited number of training batches.
We also found that although in general the time to complete one epoch takes much longer with C N C, C N C requires fewer overall training epochs for all but the CivilComments-WILDS dataset to obtain the highest reported accuracy.
D.2.6 Hyperparameter sweeps
To fairly compare with previous methods [25, 12, 42], we use the same evaluation scheme (selecting models based on worst-group validation error), and sweep over a consistent number of hyperparameters, i.e. number of validation set queries.
We set this number for C N C to be a comparable number of queries that is reported in prior works.
We break this down into method-speciﬁc (e.g. contrastive temperature in C N C, upweighting factor in J TT ), and shared (e.g. learning rate) hyperparameter categories.
Method-speciﬁc For C N C, we tune three method-speciﬁc hyperparameters: contrastive loss temperature (Eq. 1), contrastive weight (Eq. 5), and gradient accumulation steps values as in Table 3.
[Table]
For J TT , the reported results and our CMNIST ∗ implementation are tuned over the following hyperparameters in Table 4.
[Table]
For EIIL, our CMNIST ∗ and CelebA implementations are tuned over hyperparameters reported in
Table 5.
[12] report that they allow up to 20 evaluations with different hyperparameters for Waterbirds and CivilComments-WILDS.
When using GDRO as the second stage model, they also report using the same hyperparameters as the GDRO baseline for Waterbirds.
We do the same for our evaluation on CelebA.
This amounts to primarily tuning the ﬁrst stage environment inference learning rate and number of updating steps for CMNIST ∗ and CelebA, and the penalty annealing iterations and penalty weight for the IRM second stage model for CMNIST ∗ .
[Table]
Environment Inference Learning Rate
CMNIST ∗ , CelebA {1e-1, 1e-2, 1e-3}
Environment Inference Update Steps
CMNIST ∗ , CelebA {10000, 20000}
IRM Penalty Weight
CMNIST ∗ {0.1, 10, 1000, 1e5}
IRM Penalty Annealing Iterations
CMNIST ∗ {10, 50, 80}
GDRO Group Adjustment
CelebA {0, 2, 3}
For CIM, our CMNIST ∗ implementation uses CIM + VIB, and is tuned over the β VIB parameter [2] and contrastive weighting parameter λ for CIM in Table 6.
[42] report tuning over a range of values within [1e-5, 1] for λ on the CelebA and Waterbirds datasets. [Table] [Table]
[Table]
Shared For all datasets, we use the same optimizer and momentum (if applicable) as reported in the
J TT paper.
Table 7 contains the data-speciﬁc shared hyperparameter values tried. [Table] {1e [Table]
1e
[Table]
1e
[Table]
D. C N C compute resources and training time
All experiments for CMNIST ∗ , Waterbirds, and CelebA were run on a machine with 14 CPU cores and a single NVIDIA Tesla P100 GPU.
Experiments for CivilComments-WILDS were run on an
Amazon EC2 instance with eight CPUs and one NVIDIA Tesla V100 GPU.
Regarding runtime, one limitation with the current implementation of C N C is its comparatively longer training time compared to methods such as standard ERM.
This is both a result of training an initial ERM model in the ﬁrst stage, and training another model with contrastive learning in the second stage.
In Table 8 we report both how long it takes to train the initial ERM model and long it takes to complete one contrastive training epoch on each dataset.
We observe that while in some cases training the initial ERM model is negligible, especially if we employ training with only a few epochs to prevent memorization (for Colored MNIST it takes roughly two minutes to obtain a sufﬁcient initial ERM model), it takes roughly 1.5 and 3 hours to train the high regularization initial models used for Waterbirds and CelebA. While these hurdles are shared by all methods that train an initial
ERM model, we ﬁnd that the second stage of C N C occupies the bulk of training time.
Prior work has shown that contrastive learning typically requires longer training times and converges more slowly than supervised learning [9].
We also observe this in our work.
We note however that because we sample batches based on the ERM model’s predictions, the contrastive training duration is limited by how many datapoints the initial ERM model predicts incorrectly.
In moderately sized datasets with very few datapoints in minority groups, (e.g. Waterbirds, which has roughly 4794 training points and only 56 datapoints in its smallest group), the total time it takes to train C N C is on par with ERM.
Additionally, other methods such as additional hard negative mining [36] have been shown to improve the efﬁciency of contrastive learning, and we can incorporate these components to speed up training time as well.
[Table]
E Limitations and potential negative societal impacts E.1 Limitations
One limitation of our work is that it requires training an ERM model, and then training a second contrastive model based on the ERM model’s predictions.
Thus, our approach can be more expensive in terms of training time.
This could potentially be reduced by training the ERM model for fewer epochs (as in J TT ).
Another limitation is that we rely on spurious attribute labels being available on the validation set (in order to select the best model checkpoint).
This is a standard assumption made by most methods we compare to, though Sohoni et al.
[40] estimate validation group labels as well.
E.2 Potential negative societal impacts
While our work successfully improves worst-group accuracy, this is not the end-all be-all metricother fairness-based metrics may be more suitable in different settings; misuse of metrics could lead to potential harm.
To avoid these pitfalls, it is important for practitioners to understand the limitations and tradeoffs of different metrics when applying methods such as ours.
F Additional evaluations and ablations F.1 Visualization of learned data representations
In addition to the reported representation metrics in Fig. 3, we visualize and compare the learned representations of test set samples from models trained with ERM, J TT , and C N C in Fig.
Compared to ERM models, both J TT and C N C models learn representations that better depict dependencies on the class labels.
However, especially with the Waterbirds and CelebA datasets, C N C model representations more clearly depict dependencies only on the class label, as opposed to J TT models which also show some organization by the spurious attribute still.
[Figure]
F.2 Representation metric ablation with increasingly spurious datasets
To study how the relation between representation metrics and worst-group accuracy reported in Fig.and Table 1 scales with the strength of the spurious correlation, we compute representation metrics with C N C, ERM, and J TT models trained on increasingly spurious ( ↑ p corr ) CMNIST ∗ datasets in
Fig.
We observe that with high spurious correlations, ERM fails to classify digits in the minority classes, while C N C and J TT comparably maintain high worst-group accuracy.
C N C also performs better in more spurious settings ( p corr > 0 . 95 ).
These improvements over ERM are reﬂected by drops in alignment loss (averaged over classes); C N C consistently achieves lowest such loss.
Fig.
9c shows that C N C’s learned representations maintain a more favorable balance of mutual information between the class label and spurious attribute than J TT .
While J TT models exhibit slightly higher estimated
I ( Y ; Z ) than C N C models, C N C models exhibit much lower dependence on the spurious attribute.
[Figure]
F.3 Understanding C N C’s sensitivity to Stage 1 predictions
We also study how sensitive C N C is to how closely the Stage 1 ERM model actually predicts the spurious attribute.
As J TT also relies on an initial ERM model’s predictions, we compare C N C to
J TT in this regard.
We ﬁnd that C N C is more robust to noisy ERM predictions than J TT , and that
C N C does not require perfectly inferred groups to perform well.
We ﬁrst conduct an ablation on C N C and J TT ’s worst-group and average performance in Fig.
9d with the following synthetic experiment.
On CMNIST ∗ , we start with the true spurious attribute labels as the Stage 1 “predictions".
We then gradually degrade their quality as follows: for each point, with probability p we change its assigned spurious attribute label to a different label chosen uniformly at random.
Both methods’ performance degrades as p increases and the Stage 1 “predictions” degrade.
However, C N C consistently achieves higher worst-group accuracy and smaller worst-group versus average accuracy gaps.
We also observe on other datasets that C N C does not require perfectly spurious ERM predictions to work well.
For the Waterbirds and CelebA results in Table 1, the
Stage 1 ERM model predictions align with the spurious attribute value 94 .
7% and 59 .
3% of the time respectively.
While the ERM model is far from perfect at recognizing the spurious attributes, C N C still substantially reduces the worst-group vs. average accuracy gap.
F.4 Empirical validation of C N C components
We validate the design choices of C N C through various ablations studying the effects of the individual components of our method on worst-group accuracy.
F.4.1 Importance of ERM-guided contrastive sampling
Although C N C relies on an initial trained ERM model’s predictions, can we still improve worstgroup accuracy without this step and with supervised contrastive learning alone, i.e. by sampling positives uniform randomly from all datapoints with the same label as the anchor?
In Table 9, we show that this “vanilla” contrastive learning implementation competes with ERM on Waterbirds and outperforms ERM on CelebA in worst-group accuracy, while achieving similar average accuracy (standard deviation with three seeds recorded in parenthesis).
However, we still achieve substantially higher worst-group accuracy with our full method.
This supports that using initial ERM predictions to sample “hard comparisons” helps obtain best observed worst-group accuracy.
We conjecture this is the case because positive samples with the ERM model are not only the same class as anchors, but different enough such that an initial trained ERM model classiﬁed them differently.
Comparing these against negatives which are also presumably more similar to the anchors than in the no-ERM setup lends additional contrastive learning signal to not learn dependencies on spurious differences.
F.4.2 Additional design choice ablations
To validate the additional algorithmic components of C N C, we report how C N C performs on the
Waterbirds dataset when modifying the individual design components.
We use the same hyperparameters as in the main results, and report accuracies as the average over three training runs for
[Table] [Table]
the following ablations.
Table 10 summarizes that across these design ablations, default C N C as presented consistently outperforms these alternative implementations.
No projection head.
We incorporate a nonlinear projection head as is typical in prior contrastive learning works [9], that maps the encoder output to lower-dimensional representations (from 2048 to 128 in our case).
We then update the encoder layers and the projection head jointly by computing the contrastive loss on the projection head’s output, still passing the encoder layer’s direct outputs to the classiﬁer to compute the cross-entropy loss.
We note that using the projection head decreases worst-group accuracy substantially.
We reason that as previously discussed, while using the projection head in prior work can allow the model to retain more information in its actual hidden layers [9], in our case to remove dependencies on spurious attributes we actually want to encourage learning invariant representations when we model the differences between anchor and positive datapoints as due to spurious attributes.
Two-sided contrastive batches.
Instead of “two-sided” contrasting where we allow both sampled anchors and positives to take on the anchor role, for each batch we only compute contrastive updates by comparing original positives and negatives with the original anchor.
When keeping everything else the same, we ﬁnd that just doing these one-sided comparisons also leads to a drop in performance for worst-group accuracy.
This suggests that the increased number of comparisons and training setup where we swap the roles of anchors and positives of the two-sided batches introduces greater contrastive learning signal.
Joint training of encoder and classiﬁer layers.
Instead of training the full model jointly, we ﬁrst only train the encoder layers with the contrastive loss in C N C, before freezing these layers and ﬁnetuning the classiﬁer layers with the cross-entropy loss.
With this implementation, we also obtain noticeable drop in performance.
While we leave further analysis for the joint cross-entropy and contrastive optimization for future work, one conjecture is that the cross-entropy loss may aid in learning separable representations while also training the full model to keep the average error small.
From our theory, the contrastive loss can help bound the gap between worst-group and average error, and so to improve worst-group performance it may make sense to also try to minimize average error in the same parameter update.
This also follows prior work, where updating the entire model and ﬁnetuning all model parameters instead of freezing the encoder layers leads to higher accuracy [9].
However, we found that with an initial encoder-only training stage, if we did not freeze the trained layers the ﬁne-tuning on a dataset with spurious correlations would “revert” the contrastive training, resulting in a large gap between worst-group and average error similar to ERM.
F.5 Additional results
In Table 11, we provide standard deviations for C N C, as well as GDRO [37] and J TT [25] (both of which we re-ran ourselves using the publicly available source code) and G EORGE [40] (from their paper).
For GDRO and J TT , we use the reported hyperparameters in their original works respectively, but note that the J TT numbers here are slightly different than those reported in [26] (and thus Table 1).
[Table]
F.6 Additional GradCAM visualizations
On the next two pages [Table] we include additional GradCAM visualizations depicting saliency maps for samples from each group in the Waterbirds and CelebA datasets.
Warmer colors denote higher saliency [Table] suggesting that the model considered these pixels more important in making the ﬁnal classiﬁcation as measured by gradient activations.
For both datasets
[Table] we compare maps from models trained with ERM, the next most competitive method for worst-group accuracy J TT [Table] and C N C. C N C models most consistently measure highest saliency with pixels directly associated with class labels and not spurious attributes.
G Code
Anonymized code for C N C is available at this link.
gcam-g1_correct-190556-con
L a ndb i r d ,
L a nd b ac kg r ound
L a ndb i r d ,
W a t e r b a c kg r ound
W a t e r b i r d ,
L a nd b ac kg r ound
W a t e r b i r d ,
W a t e r b ac kg r ound
Input
ERM
JTT
CNC (Ours) g c a m - [-g
 _-]g1_ c o rr e c t - -
[Table]
on
[Table]
I npu t
E R M
J TT
C N C ( O u r s )
I npu t
E R M
J TT
C N C ( O u r s )
F i gu r e 11 :
A dd iti on a l G r a d
C A
M v
i s u a li za ti on s f o r t
h e C e l e b A d a t a s e t .
F o r m od e l s t r a i n e d w i t
h E
R M ,
J TT , a nd
C N
C , w e u s e
G r a d
C A
M t o v
i s u a li
ze t h e “ s a li e n t ” ob s e r v e d f ea t u r e s u s e d t o c l a ss
i f y w h e t
h
e
r a ce l e b r it
y
h
a s b l ond ( e ) h a i r .
E R
M m od e l s i n t
e
r e s ti ng a l s o “ i gno r e” t
h e ac t u a l h a i r p
i x
e l s i n f a vo r o f o t
h e r p
i x e l s , p r e s u m a b l y a ss o c
i a t e d w it h t
h e s pu
r
i
ou s g e nd e r a tt r i bu t e .
I n c on t r a s t , G r a d
C A
M s f o r J
TT a nd m od e l s u s u a ll y d e p
i
c
t h i gh e r s a li e n ce
f
o r r
e
g
i on s t h a t a t l ea s t i n c l ud e h
a i r p i x e l s .
C
N C m od e l s m o s t c on
s i s t e n tl y do s o .