Title: Towards Last-layer Retraining for Group Robustness with Fewer Annotations
Venue: NeurIPS 2023
Accepted: yes
Keywords: spurious correlations

group robustness

last-layer retraining

distribution shift
Abstract: Empirical risk minimization (ERM) of neural networks is prone to over-reliance on spurious correlations and poor generalization on minority groups. The recent deep feature reweighting (DFR) technique achieves state-of-the-art group robustness via simple last-layer retraining, but it requires held-out group and class annotations to construct a group-balanced reweighting dataset. In this work, we examine this impractical requirement and find that last-layer retraining can be surprisingly effective with no group annotations (other than for model selection) and only a handful of class annotations. We first show that last-layer retraining can greatly improve worst-group accuracy even when the reweighting dataset has only a small proportion of worst-group data. This implies a "free lunch" where holding out a subset of training data to retrain the last layer can substantially outperform ERM on the entire dataset with no additional data, annotations, or computation for training. To further improve group robustness, we introduce a lightweight method called selective last-layer finetuning (SELF), which constructs the reweighting dataset using misclassifications or disagreements. Our experiments present the first evidence that model disagreement upsamples worst-group data, enabling SELF to nearly match DFR on four well-established benchmarks across vision and language tasks with no group annotations and less than 3% of the held-out class annotations.

Towards Last-layer Retraining for Group Robustness with Fewer Annotations

Abstract

Empirical risk minimization (ERM) of neural networks is prone to over-reliance on spurious correlations and poor generalization on minority groups. The recent deep feature reweighting (DFR) technique [33] achieves state-of-the-art group robustness via simple last-layer retraining, but it requires held-out group and class annotations to construct a group-balanced reweighting dataset. In this work, we examine this impractical requirement and find that last-layer retraining can be surprisingly effective with no group annotations (other than for model selection) and only a handful of class annotations. We first show that last-layer retraining can greatly improve worst-group accuracy even when the reweighting dataset has only a small proportion of worst-group data. This implies a 'free lunch' where holding out a subset of training data to retrain the last layer can substantially outperform ERM on the entire dataset with no additional data or annotations. To further improve group robustness, we introduce a lightweight method called selective last-layer finetuning ( SELF ), which constructs the reweighting dataset using misclassifications or disagreements. Our empirical and theoretical results present the first evidence that model disagreement upsamples worst-group data, enabling SELF to nearly match DFR on four well-established benchmarks across vision and language tasks with no group annotations and less than 3% of the held-out class annotations. Our code is available at https://github.com/tmlabonte/last-layer-retraining .

1 Introduction

Classification tasks in machine learning often suffer from spurious correlations : patterns which are predictive of the target class in the training dataset but irrelevant to the true classification function. These spurious correlations, often in conjunction with the target class, create minority groups which are underrepresented in the training dataset. For example, in the task of classifying cows and camels, the training dataset may be biased so that a desert background is spuriously correlated with the camel class, creating a minority group of camels on grass backgrounds [5]. Beyond this simple scenario, spurious correlations have been observed in high-consequence applications including medicine [79], justice [9], and facial recognition [42].

Neural networks trained via the standard procedure of empirical risk minimization (ERM) [68], which minimizes the average training loss, tend to overfit to spurious correlations and generalize poorly on minority groups [18]. Even worse, it is possible for ERM models to rely exclusively on the spurious feature and incur minority group performance that is no better than random guessing [60]. Therefore, maximizing the model's group robustness , quantified by its worst accuracy on any group, is a desirable objective in the presence of spurious correlations [58].

In contrast to more generic distribution shift settings ( e.g., domain generalization [35]), the presence of spurious correlations enables the improvement of group robustness merely by addressing model bias (without collecting additional minority group data). The recently proposed deep feature reweighting (DFR) [33, 28] technique efficiently corrects model bias by retraining the last layer of the neural network, a simple procedure which achieves state-of-the-art group robustness. The key hypothesis underlying DFR is that ERM models which overfit to spurious correlations still learn core features that correlate with the ground-truth label on all groups, but they perform poorly because they overweight the spurious features in the last layer. Ostensibly, retraining the last layer on a group-balanced reweighting dataset would then upweight the core features and improve worst-group accuracy.

DFR compares favorably to existing methods such as group distributionally robust optimization (DRO) [58], which requires group annotations for the entire training dataset. However, DFR still necessitates a smaller reweighting dataset with group and class annotations to achieve maximal performance [33]. This requirement limits its practical application, as the groups are often unknown ahead of time or difficult to annotate ( e.g., due to financial, privacy, or fairness concerns).

Our contributions

In this paper, we present a comprehensive examination of the performance of last-layer retraining in the absence of group and class annotations on four well-established benchmarks for group robustness across vision and language tasks. 1 We first investigate the necessity of the reweighting dataset being balanced across groups, and we show that last-layer retraining can substantially improve worst-group accuracy even when the reweighting dataset has only a small proportion of worst-group data . Based on this observation, we propose class-balanced last-layer retraining as a simple but strong baseline for group robustness without group annotations. We show that, on average over the four datasets, this method achieves 94% of DFR worst-group accuracy compared to 76% without class balancing.

The strong performance of class-balanced last-layer retraining reveals a 'free lunch' with practical ramifications. While class-balanced ERM was recently proposed by [27] as a competitive baseline for worst-group accuracy, we show that instead of using the entire training dataset for ERM, dependence on spurious correlations can be reduced by randomly splitting the training data in two, then performing ERMtraining on the first split and class-balanced last-layer retraining on the second. Our experiments indicate that this technique can improve worst-group accuracy by up to 17% over class-balanced ERM on the original dataset using no additional data or annotations (even for model selection) - a surprising and unexplained result given that the two splits have equally drastic group imbalance.

While retraining the last layer on a class-balanced held-out dataset can be effective, it is inferior to DFR when group imbalance is large. To close this gap, we propose selective last-layer finetuning ( SELF ), which selects a small, more group-balanced reweighting dataset and finetunes the last layer instead of retraining. We implement SELF using points that are either misclassified or disagree in their predictions relative to a regularized model. Disagreement SELF does not require class annotations for the entire held-out set (only for the disagreements). We show that disagreement SELF between the ERM and early-stopped models performs the best in general, surpassing class-balanced last-layer retraining by up to 12% worst-group accuracy with less than 3% of the held-out class annotations.

Overall, our work shows benefits of last-layer retraining well beyond a group-balanced held-out dataset. Our results call for further investigation of the DFR hypothesis : while group balance is the most important factor in DFR performance, we show that a significant gain is solely due to class balancing, and the performance discrepancy between misclassification SELF and disagreement SELF suggests that worst-group accuracy may be affected by characteristics of the reweighting dataset other than group balance. Our main results are summarized and compared to previous methods in Table 1.

2 Related work

This work subsumes and improves our two previous workshop papers: [38] broadly covers Section 4, while [37] represents preliminary investigation into Section 5, which we have substantially updated with new methodology, evaluation, and theoretical analysis for this version.

Spurious correlations. The performance of empirical risk minimization (ERM) in the presence of spurious correlations has been extensively studied [18]. In vision, ERM models are widely known to rely on spurious attributes like background [58, 75], texture [17], and secondary objects [56, 61, 62] to perform classification. In language, ERM models often utilize syntactic or statistical heuristics as a substitute for semantic understanding [20, 50, 46]. This behavior can lead to bias against demographic minorities [25, 6, 67, 23, 8] or failure in high-consequence applications [42, 9, 79, 51].

Robustness and group annotations. If group annotations are available in the training dataset, group distributionally robust optimization (DRO) [58] can improve robustness by minimizing the worst-group loss, while other techniques learn invariant or diverse features [1, 19, 80, 76]. Methods which use only partial group annotations include deep feature reweighting (DFR) [33, 28], which retrains the last layer on a group-balanced held-out set, and spread spurious attribute [48], which performs DRO with group pseudo-labels. Recently, more lightweight methods that only adjust the model predictions using a group annotated held-out set have also been proposed [70]. However, since the groups are often unknown ahead of time or difficult to annotate in practice, there has been significant interest in methods which do not utilize group annotations except for model selection. The bulk of these techniques utilize auxiliary models to pseudo-label the minority group or spurious feature [63, 47, 77, 11, 31, 66, 32, 64, 81]; notably, just train twice (JTT) [41] upweights samples misclassified by an early-stopped model. Other techniques reweight or subsample the classes [27] or train with robust losses and regularization [54, 78].

Unlike DFR, our methods utilize no held-out group annotations, and unlike JTT, we do not have to 'train twice', only finetune the last layer. We also show that SELF performs best using disagreements between the ERM and early-stopped models instead of misclassifications as in JTT. Moreover, while JTT assumes the early-stopped model has low worst-group accuracy which improves during training, SELF performs well even when the early-stopped model has high worst-group accuracy which decreases during training - substantially improving performance on datasets such as CivilComments. The concurrent work of Qiu et al. [55] proposed a similar method to our Section 4.2; while they use a tunable loss to upweight misclassified samples, our method shows that similar results can be achieved with no hyperparameter tuning (and therefore no group annotations for the validation set). Finally, while our results partially corroborate the findings of Idrissi et al. [27] that class balancing during ERM is effective for group robustness, we observe that class-balanced last-layer retraining renders ERM class balancing optional. We compare our results with previous methods in Table 1.

Generalization via disagreement. Disagreement-based active learning for improving indistribution generalization has been well-studied since before the deep learning era [10, 4, 21]. More recent research has utilized disagreements between SGD runs to predict in-distribution gen-

eralization [30], as well as between model classes ( e.g., CNNs and Transformers) to predict out-ofdistribution generalization [3]. Two works that are concurrent to ours, diversity by disagreement [52] and diversify and disambiguate [40], are methods for generalization under distribution shift which maximize the disagreement between multiple predictors to learn a diverse ensemble. Compared to our work, these methods are optimized for training datasets which exhibit a complete correlation ( i.e., contain no minority group data) and they underperform in the less extreme spurious correlation setting we study.

3 Preliminaries

Setting. We consider classification tasks with input domain X and target classes Y . We assume S is a set of spurious features such that each sample x ∈ X is associated with exactly one feature s ∈ S . In conjunction with the target class, the spurious features partition the dataset into groups G = Y ×S . While the groups may be heavily imbalanced in the training distribution, we desire a model which is invariant to the spurious feature and thus has roughly uniform performance over G . Therefore, we evaluate worst-group accuracy (WGA), i.e., the minimum accuracy among all groups [58].

We will often refer to datasets and models as group-balanced or class-balanced , meaning that in expectation, the dataset is composed of an equal number of samples from groups in G or classes in Y , respectively. This balance can be achieved by training on a subset with equal data from each group/class, or sampling from the data so that each minibatch is balanced in expectation [27]. To make the latter more concrete, for group balancing we first sample s ∼ Unif ( S ) , then sample x ∼ ˆ p ( ·| s ) where ˆ p is the training distribution; class balancing is the same with Y instead of S . We use the minibatch sampling approach for both class-balanced ERM and class-balanced last-layer retraining, and we provide a comparison with the subset method in Appendix A.

Deep feature reweighting. The recently proposed deep feature reweighting (DFR) [33, 28] method achieves state-of-the-art WGA by performing ERM on the training dataset, then retraining the last layer of the neural network on a group-balanced held-out dataset, called the reweighting dataset . In the original implementation, half the validation set is used to construct the reweighting dataset: all data from the smallest group is included and the other groups are randomly downsampled to that size.

Then, the feature embeddings ( i.e., the outputs of the penultimate layer) of the reweighting dataset are pre-computed and used to train a logistic regression model with explicit glyph[lscript] 1 -regularization. The results are averaged over 10 randomly subsampled reweighting datasets, and a hyperparameter search is performed over glyph[lscript] 1 regularization strength on the other half of the validation set.

To emphasize practicality and efficiency, our implementation of DFR has some differences from the original. (i) Instead of logistic regression, we train the last layer on the reweighting dataset via minibatch optimization using SGD and AdamW [43] for the vision and language tasks, respectively. This fits well into standard training pipelines and avoids pre-computing the feature embeddings and writing them to disk, which can be slow and memory-intensive. (ii) To reduce the number of hyperparameters, we use a fixed-value glyph[lscript] 2 regularization instead of searching over glyph[lscript] 1 regularization strength. We observed similar performance for glyph[lscript] 1 and glyph[lscript] 2 regularization, which we believe is because glyph[lscript] 1 -regularized gradients do not induce sparsity [39]. (iii) We sample uniformly at random from the groups in the held-out dataset (to get group balanced minibatches) instead of averaging over group-balanced subsets of the data [27]. We compare the implementations in detail in Appendix A.

Datasets and models. We study four datasets which are well-established as benchmarks for group robustness across vision and language tasks, detailed in Table 2 and summarized below.

Waterbirds [71, 69, 58] is an image classification dataset where the task is to predict whether a bird is a landbird or a waterbird. The spurious feature is the image background: more landbirds are present on land backgrounds than waterbirds, and vice versa. 2

CelebA [42, 58] is an image classification dataset where the task is to predict whether a person is blond or not. The spurious feature is gender, with 16 × more blond women than blond men in the training set.

CivilComments [7, 35] is a text classification dataset where the task is to predict whether a comment is toxic or not. The spurious feature is the presence of one of the following categories: male, female, LGBT, black, white, Christian, Muslim, or other religion. 3 More toxic comments contain one of these categories than non-toxic comments, and vice versa.

MultiNLI [73, 58] is a text classification dataset where the task is to predict whether a pair of sentences is a contradiction, entailment, or neither. The spurious feature is a negation in the second sentence - more contradictions have this property than entailments or neutral pairs.

Importantly, Waterbirds is the only dataset that has a distribution shift - its validation and test datasets are group-balanced conditioned on the classes, although still class-imbalanced. 4 As a result, methods which use the validation set for training can improve performance without explicit group balancing.

We utilize a ResNet-50 [24] pretrained on ImageNet-1K [57] for Waterbirds and CelebA, and a BERT [14] model pretrained on Book Corpus [83] and English Wikipedia for CivilComments and MultiNLI. Following previous work, we use half the validation set for feature reweighting [33, 28] and half for model selection with group annotations [58, 41, 33, 48, 28]. We run each experiment on three random seeds and do not utilize explicit early stopping except as part of SELF (see Section 5). See Appendix D for further details on the training procedure.

4 Class-balanced last-layer retraining

In this section, we investigate the necessity of a group-balanced reweighting dataset for DFR and show that class-balanced last-layer retraining is a simple but strong baseline for group robustness without group annotations. To enable a fair comparison with our implementation of DFR (see Section 3), class-balanced last-layer retraining follows the same training procedure, except the reweighting dataset is constructed by sampling uniformly over the classes Y instead of the groups G .

0

20

40

60

80

100

% of worst-group data included

50

60

70

80

90

Worst-group test accuracy

Waterbirds

CelebA

CivilComments

ERM Baselines

(a) Worst-group test accuracy

0

20

40

60

80

100

% of worst-group data included

0

20

40

60

80

100

Relative % of max WGA increase over ERM

Waterbirds

CelebA

CivilComments

ERM Baselines

(b) Relative % of max WGA increase over ERM

4.1 An ablation on the proportion of worst-group data

In this section, we perform an ablation on the percentage of worst-group data in the reweighting dataset and show that class-balanced last-layer retraining can substantially improve WGA even when the reweighting dataset has only a small proportion of worst group data. For our ablation, we choose the worst groups based on the performance of ERM, and if two groups have similarly poor performance, we vary both. The worst groups for Waterbirds are landbirds on water backgrounds and waterbirds on land backgrounds; for CelebA blond men; for CivilComments toxic comments containing identity categories; and for MultiNLI entailments and neutral pairs containing negations.

For these experiments, we begin with the DFR reweighting dataset, i.e., a random group-balanced subset of the held-out dataset, where the size of each group is the minimum of the worst-group size and half the size of any other group. 5 We define this dataset to include 100% of the worst-group data. We then reduce the percentage of worst-group data, while correspondingly increasing the percentage of data from the same class without the spurious feature. For example, the DFR reweighting dataset on CelebA has 92 points from each group, so reducing the worst group to 25% results in 92 non-blond females, 92 non-blond males, 161 blond females, and 23 blond males. The total data is kept constant.

Waterbirds

CelebA

CivilComments

50

55

60

65

70

75

80

85

90

Worst-group test accuracy

Class-balanced (CB) ERM

CB ERM and CB last-layer retraining

(a) Training dataset only

Waterbirds

CelebA

CivilComments

50

55

60

65

70

75

80

85

90

Worst-group test accuracy

Class-balanced (CB) ERM

CB ERM and CB last-layer retraining

(b) Combined training and held-out datasets

The results of this ablation are displayed in Figure 1. While a smooth increase in WGA with increasing worst-group data is expected, the extent of the early increase is surprising: over the four datasets, an average of 67% of the increase in WGA over class-unbalanced ERM is obtained by the first 25% of worst-group data. In particular, CelebA and CivilComments - the most class-imbalanced datasets we study - experience significant improvement even at low percentages of worst-group data. This phenomenon suggests that, while group balancing is still important for best results, class balancing is a major factor in the performance of DFR on these two datasets.

Furthermore, class-balanced last-layer retraining can improve worst-group accuracy even when the held-out dataset has similar group imbalance as the training dataset ( i.e., at the stars 8 in Figure 1). Based on this observation, we propose that class-balanced last-layer retraining on the entire held-out dataset can be a simple but strong baseline for group robustness without group annotations. Table 3 details the results; on average over the four datasets, class-balanced last-layer retraining achieves 94% of DFR performance, compared to 76% without class balancing. However, it still trails DFR by a significant amount on CelebA and MultiNLI - the most group-imbalanced datasets we study. We improve these results with our selective last-layer finetuning ( SELF ) method described in Section 5.

Moreover, our experiments in Figure 4 (deferred to Appendix A) indicate that class-balanced last layer retraining has similar performance regardless of whether it is initialized with class-unbalanced or class-balanced ERM features. Contrasting with Idrissi et al. [27], our results suggest that class balancing in the ERM stage is optional. This result has practical relevance for expensive models pre-trained without class balancing ( e.g., large language models), as the benefits of class balancing on downstream tasks can be reaped by simply retraining the last layer instead of training a new model.

4.2 A 'free lunch' in group robustness

Motivated by the promising results of Section 4.1, where we performed class-balanced last-layer retraining on a fixed held-out set, we now ask how can we best utilize a realistic training dataset? In particular, practical applications often work with a predetermined data, annotation, and compute budget. Within this budget, and with no explicit held-out dataset, would one achieve better group robustness by using the entire dataset for ERM or by holding out a subset for last-layer retraining?

We investigate this question on our four benchmark datasets by randomly splitting the initial dataset into two, then performing class-balanced ERM training on the first split ( 95% of the data) and class-balanced last-layer retraining on the second split ( 5% of the data). Figure 2a illustrates the results of our experiments on the training dataset and Figure 2b on the combined training and held-out datasets. Since the quantity of data is higher in the case of combining the training and held-out datasets, we expect all numbers to be higher in Figure 2b compared to Figure 2a (especially on

Waterbirds, which has a more group-balanced validation set). We perform no hyperparameter tuning for these experiments, and therefore we do not utilize group annotations even for model selection .

Figure 2 indicates that splitting the dataset and performing last-layer retraining substantially improves worst-group accuracy on Waterbirds, CelebA, and CivilComments. It decreases performance on MultiNLI, which is the only dataset where adding held-out data from the same distribution significantly increases ERM worst-group accuracy compared to DFR. Specifically, class-balanced ERM achieves 67 . 4 ± 2 . 4 %WGAonthe training dataset and 69 . 2 ± 1 . 6 %on the combined training and held-out datasets, while our DFR implementation achieves 70 . 8 ± 0 . 8 %. Therefore, we hypothesize that last-layer retraining on the second split can improve group robustness only if there is enough data for ERM to perform near-optimally on the first split, i.e., if the performance of ERM on the first split is limited by dataset bias rather than sample variance.

Based on this hypothesis, our answer to the posed question is: if ERM performance is stable when holding out 5% of data, perform last-layer retraining on the held-out dataset instead of ERM on the initial dataset. We call this technique a 'free lunch' because it improves worst-group accuracy with no additional data or annotations beyond ERM (including for model selection). In particular, we utilize less data for ERM training and less compute due to the efficient nature of last-layer retraining. Therefore, we believe this method is especially relevant to practitioners, and it can be easily implemented with little change to data processing or model training workflows.

Acritical remaining question is why last-layer retraining improves group robustness; since the training and held-out datasets have equally drastic group imbalance, it is counterintuitive that reducing the quantity of data used for ERM and performing last-layer retraining would increase worst-group accuracy. In some cases, as seen on MultiNLI in Figure 2, training an ERM model with less data can be detrimental - but on the other three datasets, last-layer retraining substantially improves over ERM. We leave it to future empirical and theoretical work to better understand this phenomenon.

5 Selective last-layer finetuning

While class-balanced last-layer retraining can improve worst-group accuracy without group annotations, its performance is still inferior to DFR, particularly on the highly group-imbalanced CelebA and MultiNLI datasets. In these cases, training on the entire held-out set can be detrimental; instead, we show that constructing the reweighting dataset by detecting and upsampling worst-group data is more effective. We propose selective last-layer finetuning ( SELF ), which uses an auxiliary model to select a small, more group-balanced reweighting dataset, then finetunes the ERM last layer instead of retraining entirely to avoid overfitting on this smaller dataset.

SELF does not use group annotations for training and can be implemented even when class annotations are unavailable in the held-out dataset, making it applicable in settings not captured by current techniques ( e.g., adaptation to an unlabeled target domain [45]). If class annotations are not available, we select points which are disagreed upon by ERM and a regularized model, then request their labels. If class annotations are available, we can still use disagreement, or alternatively, we select points which are misclassified by the ERM or early-stopped model. Our experiments indicate that disagreement SELF between the ERM and early-stopped models performs the best overall, increasing the worst-group accuracy of class-balanced last-layer retraining to near-DFR levels while requesting less than 3% of the held-out class annotations. Our results are detailed in Table 4.

We formalize SELF as follows. Let f, g : X → R |Y| be models with logit outputs, X ⊆ X be a held-out dataset, c : R 2 |Y| → R be a cost function, and n be an integer. SELF constructs the reweighting set D by greedily selecting the n points in X whose outputs incur the highest cost, i.e.,

Then, SELF requests class annotations and performs class-balanced finetuning ( i.e., starting from the ERM weights) of the last layer of the ERM model on D , using the same implementation as last-layer retraining (see Section 4). We study the following four variants of SELF :

Misclassification : f is an ERM model and g is the labeling function, i.e., , for a datapoint x with label y , the output g ( x ) is the one-hot encoded vector for y . The cost function c is cross entropy loss.

Early-stop misclassification : 6 f is an early-stopped ERM model and g is the labeling function. The cost function c is cross entropy loss.

Dropout disagreement : f is an ERM model and g is the same model where inference is run with nodes randomly dropped out in the last layer [65]. The cost function c is KL divergence with softmax.

Early-stop disagreement : f is an ERM model and g is an early-stopped ERM model. The cost function c is KL divergence with softmax.

We also experiment with a pure random baseline, referred as Random SELF in Table 4, where D consists of n random points from X . While Random SELF suffers high variance as expected, perhaps surprisingly its performance is still competitive on average. The reweighting dataset size n is a hyperparameter, and for disagreement SELF it corresponds to the quantity of class annotations requested. Please see Appendices B and D for hyperparameter details and ablation studies.

5.1 Analysis of SELF performance

To quantify the extent of SELF 's upsampling, we plot the percentage of the reweighting dataset consisting of worst-group data in Figure 3. In particular, we present to the best of our knowledge the first evidence that model disagreement effectively upsamples worst-group data, an observation which may be of independent interest. With that said, Figure 3 does not formally establish that better group balance has a strong correlation with WGA - rather, it suggests an additional subtlety in the DFR hypothesis , since group balance alone does not fully explain the performance of last-layer retraining. In addition to the balance of the reweighting dataset, it is likely that characteristics of the specific data selected also contribute to SELF results. Additionally, the competitive performance of pure random SELF in Table 4 further questions the importance of group balancing in DFR.

The importance of which data are selected could explain why disagreement SELF often outperforms misclassification SELF despite having access to less information. While misclassification selects the most difficult or noisiest points, disagreement selects the most uncertain points. For example, dropout models approximate a theoretically justified uncertainty metric [16] which is likely to be higher on worst-group data, and early-stopped models tend to fit simple patterns first [2, 41] including the majority group. Training on the most uncertain data is a key tenet of active learning [13, 59, 12], which rationalizes the performance of disagreement SELF and provides motivation for further investigation of why last-layer retraining improves group robustness.

Waterbirds

CelebA

CivilComments

0

20

40

60

% of worst-group data

Baseline

Balanced

Misclassification

ES misclassification

Dropout disagreement

ES disagreement

A remaining question is why the performance of disagreement SELF is so strong on CivilComments compared to misclassification-based methods such as JTT [41] and ES misclassification SELF . We show in Figure 6 (deferred to Appendix B) that, contrary to the assumptions made by JTT and other early-stop misclassification methods, the worst-group accuracy decreases with training on CivilComments. This highlights a potential advantage of disagreement methods over misclassification methods, as disagreement is justified regardless of whether the regularized or ERM model has a greater dependence on spurious features . Moreover, while misclassification methods do indeed upsample the minority group in Figure 3, we show in Table 10 (deferred to Appendix B) that the held-out set training accuracy of misclassification SELF tends to be very low - down to 0% on MultiNLI - evidence that misclassifications are too difficult to learn without modifying the features.

Theoretical proof-of-concept. The observations of Section 5.1 raise the question of whether disagreement SELF can ever be shown to provably upsample minority group points . We provide a simple proof-of-concept for SELF by considering the last layer to be a linear model (which could be under- or overparameterized) with core, spurious and junk features. We show that the disagreement between the regularized model and the ERM model (measured by total variation distance between the predicted distributions) is provably higher on minority examples than majority examples, regardless of model dependence on the spurious feature. In particular, this shows provable benefits of disagreement SELF even in situations where the early-stopped model has higher worst-group accuracy than the convergent model, which may not be captured by related methods in the literature [41]. Our detailed setup, assumptions, and main theoretical result (Theorem 1) are stated in detail in Appendix C.

6 Conclusion

In this paper, we presented a comprehensive examination of the performance of last-layer retraining in the absence of group and class annotations. We showed that class-balanced last-layer retraining is a simple but strong baseline for group robustness, and that holding out a subset of the training data to retrain the last layer can substantially improve worst-group accuracy. We then proposed selective last-layer finetuning ( SELF ), whose early-stop disagreement version improves performance to near-DFR levels with no group annotations and less than 3% of the held-out class annotations.

Our work has generated several open questions which could prove fruitful for further research. First, why does last-layer retraining on a held-out split of the training dataset improve group robustness (see Section 4.2)? Second, to what extent does the precise data selected matter for reweighting, e.g., which of disagreement SELF and misclassification SELF would perform better when hyperparameters are set to equalize group balance? Third, what is the optimal strategy in general for obtaining the regularized model in disagreement SELF (dropout, early-stopping, or a different technique) and why?

Acknowledgments. We thank Google Cloud for the gift of compute credits, Jacob Abernethy and Eva Dyer for the additional compute assistance, and anonymous reviewers for their helpful feedback. We also thank Harikrishna Narasimhan for reading an earlier version of the manuscript and providing thoughtful comments and Seokhyeon Jeong for catching an error in an earlier version of the code. T.L. acknowledges support from the DoD NDSEG Fellowship. V.M. acknowledges support from the NSF (through CAREER award CCF-2239151 and award IIS-2212182), an Adobe Data Science Research Award, an Amazon Research Award, and a Google Research Collabs Award.

References

Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint 1907.02893 , 2019. Cited on page 3.

Devansh Arpit, Stanisław Jastrz˛ ebski, Nicolas Ballas, David Krueger, Emmanuel Bengio, Maxinder S. Kanwal, Tegan Maharaj, Asja Fischer, Aaron Courville, Yoshua Bengio, and Simon Lacoste-Julien. A closer look at memorization in deep networks. In International Conference on Machine Learning (ICML) , 2017. Cited on page 9.

Christina Baek, Yiding Jiang, Aditi Raghunathan, and Zico Kolter. Agreement-on-the-line: Predicting the performance of neural networks under distribution shift. In Conference on Neural Information Processing Systems (NeurIPS) , 2022. Cited on page 4.

Maria-Florina Balcan, Alina Beygelzimer, and John Langford. Agnostic active learning. In International Conference on Machine Learning (ICML) , 2006. Cited on page 3.

Sara Beery, Grant van Horn, and Pietro Perona. Recognition in terra incognita. In European Conference on Computer Vision (ECCV) , 2018. Cited on page 1.

Su Lin Blodgett, Lisa Green, and Brendan O'Connor. Demographic dialectal variation in social media: A case study of African-American English. In Empirical Methods in Natural Language Processing (EMNLP) , 2016. Cited on page 3.

Daniel Borkan, Lucas Dixon, Jeffrey Sorenson, Nithium Thain, and Lucy Vasserman. Nuanced metrics for measuring unintended bias with real data for text classification. In World Wide Web (WWW) , 2019. Cited on page 5.

Joy Buolamwini and Timnit Gebru. Gender shades: Intersectional accuracy disparities in commercial gender classification. In Conference on Fairness, Accountability, and Transparency in Machine Learning (FATML) , 2018. Cited on page 3.

Alexandra Chouldechova. Fair prediction with disparate impact: A study of bias in recidivism prediction instruments. In Conference on Fairness, Accountability, and Transparency in Machine Learning (FATML) , 2016. Cited on pages 1 and 3.

David Cohn, Les Atlas, and Richard Ladner. Improving generalization with active learning. Machine Learning , 15(2):201-221, 1994. Cited on page 3.

Elliot Creager, Jörn-Henrik Jacobsen, and Richard Zemel. Environment inference for invariant learning. In International Conference on Machine Learning (ICML) , 2021. Cited on page 3.

Aron Culotta and Andrew McCallum. Reducing labeling effort for structured prediction tasks. In Conference on Artificial Intelligence (AAAI) , 2005. Cited on page 9.

Ido Dagan and Sean P. Engelson. Committee-based sampling for training probabilistic classifiers. In International Conference on Machine Learning (ICML) , 1995. Cited on page 9.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In North American Association for Computational Linguistics (NAACL) , 2019. Cited on pages 5 and 25.

Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M. Roy, and Surya Ganguli. Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel. In Conference on Neural Information Processing Systems (NeurIPS) , 2020. Cited on page 23.

Yarin Gal and Zoubin Ghahramani. Dropout as a Bayesian approximation: Representing model uncertainty in deep learning. In International Conference on Machine Learning (ICML) , 2016. Cited on page 9.

Robert Geirhos, Patricia Rubisch, Claudio Michaelis, Matthias Bethge, Felix A. Wichmann, and Wieland Brendel. ImageNet-trained CNNs are biased towards texture; increasing shape bias improves accuracy and robustness. In International Conference on Learning Representations (ICLR) , 2019. Cited on page 3.

Robert Geirhos, Jörn-Henrik Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix A. Wichmann. Shortcut learning in deep neural networks. Nature Machine Intelligence , 2:665-673, 2020. Cited on pages 1 and 3.

Karan Goel, Albert Gu, Yixuan Li, and Christopher Ré. Model patching: Closing the subgroup performance gap with data augmentation. In International Conference on Learning Representations (ICLR) , 2021. Cited on page 3.

Suchin Gururangan, Swabha Swayamdipta, Omer Levy, Roy Schwartz, Samuel Bowman, and Noah A. Smith. Annotation artifacts in natural language inference data. In North American Association for Computational Linguistics (NAACL) , 2018. Cited on page 3.

Steve Hanneke. Theory of Disagreement-Based Active Learning . Foundations and Trends in Machine Learning, 2014. Cited on page 3.

Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Haldane, Jaime Fernández del Río, Mark Wiebe, Pearu Peterson, Pierre Gérard-Marchant, Kevin Sheppard, Tyler Reddy, Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant. Array programming with NumPy. Nature , 585(1):357-362, 2020. Cited on page 25.

Tatsunori B. Hashimoto, Megha Srivastava, Hongseok Namkoong, and Percy Liang. Fairness without demographics in repeated loss minimization. In International Conference on Machine Learning (ICML) , 2018. Cited on page 3.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Conference on Computer Vision and Pattern Recognition (CVPR) , 2016. Cited on pages 5 and 25.

Dirk Hovy and Anders Søgaard. Tagging performance correlates with author age. In Association for Computational Linguistics (ACL) , 2015. Cited on page 3.

John D. Hunter. Matplotlib: A 2D graphics environment. Computing in Science & Engineering , 9(3):90-95, 2007. Cited on page 25.

Badr Youbi Idrissi, Martín Arjovsky, Mohammad Pezeshki, and David Lopez-Paz. Simple data balancing achieves competitive worst-group-accuracy. In Conference on Causal Learning and Reasoning (CLeaR) , 2022. Cited on pages 2, 3, 4, 5, 7, and 25.

Pavel Izmailov, Polina Kirichenko, Nate Gruver, and Andrew Gordon Wilson. On feature learning in the presence of spurious correlations. In Conference on Neural Information Processing Systems (NeurIPS) , 2022. Cited on pages 2, 3, 4, 5, 6, 8, 18, 25, and 28.

Stanisław Jastrz˛ ebski, Maciej Szymczak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho, and Krzysztof Geras. The break-even point on optimization trajectories of deep neural networks. In International Conference on Learning Representations (ICLR) , 2020. Cited on page 23.

Yiding Jiang, Vaishnavh Nagarajan, Christina Baek, and J. Zico Kolter. Assessing generalization of SGD via disagreement. In International Conference on Learning Representations (ICLR) , 2022. Cited on page 4.

Eungyeup Kim, Jihyeon Lee, and Jaegul Choo. BiaSwap: removing dataset bias with biastailored swapping augmentation. In International Conference on Computer Vision (ICCV) , 2021. Cited on page 3.

Nayeong Kim, Sehyun Hwang, Sungsoo Ahn, Jaesik Park, and Suha Kwak. Learning debiased classifier with biased committee. In Conference on Neural Information Processing Systems (NeurIPS) , 2022. Cited on page 3.

Polina Kirichenko, Pavel Izmailov, and Andrew Gordon Wilson. Last layer re-training is sufficient for robustness to spurious correlations. In International Conference on Learning Representations (ICLR) , 2023. Cited on pages 1, 2, 3, 4, 5, 6, 8, 18, 25, and 28.

Frederic Koehler, Lijia Zhou, Danica J. Sutherland, and Nathan Srebro. Uniform convergence of interpolators: Gaussian width, norm bounds and benign overfitting. In Conference on Neural Information Processing Systems (NeurIPS) , 2021. Cited on page 23.

Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Sara Beery, Jure Leskovec, Anshul Kundaje, Emma Pierson, Sergey Levine, Chelsea Finn, and Percy Liang. WILDS: A benchmark of in-the-wild distribution shifts. In International Conference on Machine Learning (ICML) , 2021. Cited on pages 2 and 5.

Tyler LaBonte. Milkshake: Quick and extendable experimentation with classification models. http://github.com/tmlabonte/milkshake , 2023. Cited on page 25.

Tyler LaBonte, Vidya Muthukumar, and Abhishek Kumar. Dropout disagreement: A recipe for group robustness with fewer annotations. In Conference on Neural Information Processing Systems (NeurIPS) Workshop on Distribution Shifts , 2022. Cited on page 2.

Tyler LaBonte, Vidya Muthukumar, and Abhishek Kumar. Saving a split for last-layer retraining can improve group robustness without group annotations. In International Conference on Machine Learning (ICML) Workshop on Spurious Correlations, Invariance, and Stability , 2023. Cited on page 2.

John Langford, Lihong Li, and Tong Zhang. Sparse online learning via truncated gradient. Journal of Machine Learning Research (JMLR) , 10(1):777-801, 2009. Cited on page 5.

Yoonho Lee, Huaxiu Yao, and Chelsea Finn. Diversify and disambiguate: Learning from underspecified data. In International Conference on Learning Representations (ICLR) , 2023. Cited on pages 4 and 28.

Evan Zheran Liu, Behzad Haghgoo, Annie S. Chen, Aditi Raghunathan, Pang Wei Koh, Shiori Sagawa, Percy Liang, and Chelsea Finn. Just train twice: Improving group robustness without training group information. In International Conference on Machine Learning (ICML) , 2021. Cited on pages 2, 3, 5, 9, 10, 21, and 28.

Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In International Conference on Computer Vision (ICCV) , 2015. Cited on pages 1, 3, and 5.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations (ICLR) , 2019. Cited on pages 5 and 25.

TorchVision maintainers and contributors. TorchVision: PyTorch's computer vision library. GitHub , 2016. Cited on page 25.

Yishay Mansour, Mehryar Mohri, and Afshin Rostamizadeh. Domain adaptation: Learning bounds and algorithms. In Conference on Learning Theory (COLT) , 2009. Cited on page 9.

Tom McCoy, Ellie Pavlick, and Tal Linzen. Right for the wrong reasons: Diagnosing syntactic heuristics in natural language inference. In Association for Computational Linguistics (ACL) , 2019. Cited on page 3.

Junhyun Nam, Hyuntak Cha, Sungsoo Ahn, Jaeho Lee, and Jinwoo Shin. Learning from failure: Training debiased classifier from biased classifier. In Conference on Neural Information Processing Systems (NeurIPS) , 2020. Cited on page 3.

Junhyun Nam, Jaehyung Kim, Jaeho Lee, and Jinwoo Shin. Spread spurious attribute: Improving worst-group accuracy with spurious attribute estimation. In International Conference on Learning Representations (ICLR) , 2022. Cited on pages 2, 3, 5, and 28.

Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. arXiv preprint arXiv:1412.6614 , 2014. Cited on page 23.

Timothy Niven and Hung-Yu Kao. Probing neural network comprehension of natural language arguments. In Association for Computational Linguistics (ACL) , 2019. Cited on page 3.

Luke Oakden-Rayner, Jared Dunnmon, Gustavo Carneiro, and Christopher Ré. Hidden stratification causes clinically meaningful failures in machine learning for medical imaging. In Conference on Neural Information Processing Systems (NeurIPS) Workshop on Machine Learning for Health , 2019. Cited on page 3.

Matteo Pagliardini, Martin Jaggi, François Fleuret, and Sai Praneeth Karimireddy. Agree to disagree: Diversity through disagreement for better transferability. In International Conference on Learning Representations (ICLR) , 2023. Cited on pages 4 and 28.

Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in PyTorch. In Conference on Neural Information Processing Systems (NeurIPS) Workshop on Automatic Differentiation , 2017. Cited on page 25.

Mohammad Pezeshki, Sékou-Oumar Kaba, Yoshua Bengio, Aaron Courville, Doina Precup, and Guillaume Lajoie. Gradient starvation: A learning proclivity in neural networks. In Conference on Neural Information Processing Systems (NeurIPS) , 2021. Cited on page 3.

Shikai Qiu, Andres Potapczynski, Pavel Izmailov, and Andrew Gordon Wilson. Simple and fast group robustness by automatic feature reweighting. In International Conference on Machine Learning (ICML) , 2023. Cited on page 3.

Amir Rosenfeld, Richard Zemel, and John K. Tsotsos. The elephant in the room. arXiv preprint 1808.03305 , 2018. Cited on page 3.

Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet large scale visual recognition challenge. International Journal of Computer Vision (IJCV) , 115(1):211-252, 2015. Cited on pages 5 and 25.

Shiori Sagawa, Pang Wei Koh, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. In International Conference on Learning Representations (ICLR) , 2020. Cited on pages 1, 2, 3, 4, 5, 25, and 28.

Tobias Scheffer, Christian Decomain, and Stefan Wrobel. Active hidden markov models for information extraction. In International Conference on Advances in Intelligent Data (CAIDA) , 2001. Cited on page 9.

Harshay Shah, Kaustav Tamuly, Aditi Raghunathan, Prateek Jain, and Praneeth Netrapalli. The pitfalls of simplicity bias in neural networks. In Conference on Neural Information Processing Systems (NeurIPS) , 2020. Cited on page 1.

Rakshith Shetty, Bernt Schiele, and Mario Fritz. Not using the car to see the sidewalk: Quantifying and controlling the effects of context in classification and segmentation. In Conference on Computer Vision and Pattern Recognition (CVPR) , 2019. Cited on page 3.

Sahil Singla and Soheil Feizi. Salient ImageNet: how to discover spurious features in deep learning? In International Conference on Learning Representations (ICLR) , 2022. Cited on page 3.

Nimit S. Sohoni, Jared A. Dunnmon, Geoffrey Angus, Albert Gu, and Christopher Ré. No subclass left behind: Fine-grained robustness in coarse-grained classification problems. In Conference on Neural Information Processing Systems (NeurIPS) , 2020. Cited on page 3.

Nimit S. Sohoni, Maziar Sanjabi, Nicolas Ballas, Aditya Grover, Shaoliang Nie, Hamed Firooz, and Christopher Ré. BARACK: partially supervised group robustness with guarantees. In International Conference on Machine Learning (ICML) Workshop on Spurious Correlations, Invariance, and Stability , 2022. Cited on page 3.

Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: A simple way to prevent neural networks from overfitting. Journal of Machine Learning Research , 15(1):1929-1958, 2014. Cited on page 9.

Saeid Asgari Taghanaki, Aliasghar Khani, Fereshte Khani, Ali Gholami, Linh Trana, Ali Mahdavi-Amiri, and Ghassan Hamarneh. MaskTune: mitigating spurious correlations by forcing to explore. In Conference on Neural Information Processing Systems (NeurIPS) , 2022. Cited on pages 3 and 5.

Rachael Tatman. Gender and dialect bias in YouTube's automatic captions. In Association for Computational Linguistics (ACL) Workshop on Ethics in Natural Language Processing , 2017. Cited on page 3.

Vladimir Vapnik. Statistical Learning Theory . Wiley, 1998. Cited on page 1.

Catherine Wah, Steve Branson, Peter Welinder, Pietro Perona, and Serge Belongie. The CaltechUCSD birds-200-2011 dataset. Technical report, California Institute of Technology, 2011. Cited on page 5.

Jiaheng Wei, Harikrishna Narasimhan, Ehsan Amid, Wen-Sheng Chu, Yang Liu, and Abhishek Kumar. Distributionally robust post-hoc classifiers under prior shifts. In The Eleventh International Conference on Learning Representations , 2022. Cited on page 3.

Peter Welinder, Steve Branson, Takeshi Mita, Catherine Wah, Florian Schroff, Serge Belongie, and Pietro Perona. Caltech-UCSD birds 200. Technical report, California Institute of Technology, 2010. Cited on page 5.

William Falcon and the PyTorch Lightning maintainers and contributors. PyTorch Lightning. GitHub , 2019. Cited on page 25.

Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In North American Association for Computational Linguistics (NAACL) , 2018. Cited on page 5.

Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Remi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander Rush. Transformers: State-of-theart natural language processing. In Conference on Empirical Methods in Natural Language Processing (EMNLP) System Demonstrations , 2020. Cited on page 25.

Kai Yuanqing Xiao, Logan Engstrom, Andrew Ilyas, and Aleksander M ˛ adry. Noise or signal: The role of image backgrounds in object recognition. In International Conference on Learning Representations (ICLR) , 2021. Cited on page 3.

Yilun Xu, Hao He, Tianxiao Shen, and Tommi Jaakkola. Controlling directions orthogonal to a classifier. In International Conference on Learning Representations (ICLR) , 2022. Cited on page 3.

Yadollah Yaghoobzadeh, Soroush Mehri, Remi Tachet, T.J. Hazen, and Alessandro Sordoni. Increasing robustness to spurious correlations using forgettable examples. In European Association for Computational Linguistics (EACL) , 2021. Cited on page 3.

Yao-Yuan Yang, Chi-Ning Chou, and Kamalika Chaudhuri. Understanding rare spurious correlations in neural networks. arXiv preprint 2202.05189 , 2022. Cited on page 3.

John R. Zech, Marcus A. Badgeley, Manway Liu, Anthony B. Costa, Joseph J. Titano, and Eric Karl Oermann. Variable generalization performance of a deep learning model to detect pneumonia in chest radiographs: A cross-sectional study. PLoS Medicine , 15:e1002683, 2018. Cited on pages 1 and 3.

Jianyu Zhang, David Lopez-Paz, and Léon Bottou. Rich feature construction for the optimization-generalization dilemma. In International Conference on Machine Learning (ICML) , 2022. Cited on page 3.

Michael Zhang, Nimit S. Sohoni, Hongyang R. Zhang, Chelsea Finn, and Christopher Ré. Correct-n-contrast: A contrastive approach for improving robustness to spurious correlations. In International Conference on Machine Learning (ICML) , 2022. Cited on pages 3 and 5.

Lijia Zhou, Danica J. Sutherland, and Nathan Srebro. On uniform convergence and low-norm interpolation learning. In Conference on Neural Information Processing Systems (NeurIPS) , 2020. Cited on page 23.

Yukun Zhu, Ryan Kiros, Rich Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In International Conference on Computer Vision (ICCV) , 2015. Cited on pages 5 and 25.

Appendix Outline

The appendix is organized as follows. In Appendix A we include additional experiments for Sections 3 and 4. In Appendix B we include additional experiments for Section 5. In Appendix C we include the statement and proof of our theoretical proof-of-concept (Theorem 1). In Appendix D we describe in detail our training procedure and hyperparameters. Appendix E include tables to accompany the figures in the main paper. Finally, Appendix F discusses broader impacts, limitations, and compute.

A Additional experiments for Sections 3 and 4

In this section we include the detailed results of additional experiments for Sections 3 and 4.

Waterbirds

CelebA

CivilComments

50

55

60

65

70

75

80

85

90

95

Worst-group test accuracy

Retraining with UB ERM features

Retraining with CB ERM features

B Additional experiments for Section 5

In this section we include the detailed results of additional experiments for Section 5.

20

100

500

Reweighting dataset size

75

80

85

90

Worst-group accuracy (%)

CB ERM

CB last-layer retraining

DFR

Misclassification

ES misclassification

Dropout disagreement

ES disagreement

Random

(a) Waterbirds

20

100

500

Reweighting dataset size

60

65

70

75

80

85

Worst-group accuracy (%)

(b) CelebA

20

100

500

Reweighting dataset size

60

65

70

75

80

Worst-group accuracy (%)

(c) CivilComments

20

100

500

Reweighting dataset size

65

66

67

68

69

70

71

Worst-group accuracy (%)

20

40

60

80

100

Epoch

65

70

75

80

85

90

95

100

Group Accuracy (%)

landbird/land

landbird/water

waterbird/land

waterbird/water

(a) Waterbirds

5

10

15

20

Epoch

65

70

75

80

85

90

95

100

Group Accuracy (%)

non-blond/female

non-blond/male

blond/female

non-blond/male

(b) CelebA

2

4

6

8

10

Epoch

60

65

70

75

80

85

90

95

Group Accuracy (%)

neutral/no identity

neutral/identity

toxic/no identity

toxic/identity

(c) CivilComments

2

4

6

8

10

Epoch

60

65

70

75

80

85

90

95

Group Accuracy (%)

contradiction/no negation

contradiction/negation

entailment/no negation

entailment/negation

neither/no negation

neither/negation

1

2

5

10

20

50

100

% of group annotations included

0

20

40

60

80

100

Worst-group test accuracy

ES disagreement SELF

DFR (our impl.)

Class-balanced ERM

(a) Waterbirds

1

2

5

10

20

50

100

% of group annotations included

50

60

70

80

90

Worst-group test accuracy

ES disagreement SELF

DFR (our impl.)

Class-balanced ERM

(b) CelebA

1

2

5

10

20

50

100

% of group annotations included

20

30

40

50

60

70

80

90

Worst-group test accuracy

ES disagreement SELF

DFR (our impl.)

Class-balanced ERM

(c) CivilComments

1

2

5

10

20

50

100

% of group annotations included

58

60

62

64

66

68

70

72

Worst-group test accuracy

ES disagreement SELF

DFR (our impl.)

Class-balanced ERM

C Theoretical proof-of-concept for disagreement SELF

In tthis section, we describe and state our main theoretical proof-of-concept of disagreement SELF (Theorem 1). We begin with a description of our setup and assumptions.

Setup. Let X ⊆ R v be the data domain, n be the sample size, and d ≥ v, n be the number of features. For x ∈ X and weights α, β > 0 we assume that the learned predictor of both the ERM and regularized models takes the form

and assume that the label y is generated by y = sgn ( φ core ( x )) ( i.e., zero label noise). Note that this specific split of features into core, spurious, and junk categories is an extension of a common simplified setting in the literature on overparameterized linear models [82, 34], where only core and junk features are considered and only in-distribution accuracy is analyzed. We also consider the domain X to be partitioned into a minority group X min and majority group X maj such that yφ sp ( x ) < 0 for all x ∈ X min and yφ sp ( x ) > 0 for all x ∈ X maj. (By definition, yφ core ( x ) > 0 for any x .)

Assumptions on the feature learning process. We make some extra simplifying assumptions on the feature learning process that are motivated by empirical observations in the literature. First, we assume that the learned features φ core , φ sp , φ junk are identical for the regularized and ERM models, and are only weighted differently in the last layer. This assumption is well-justified for dropout disagreement SELF since we only use dropout for inference in the last layer (with frozen features), and for early-stop disagreement SELF since feature learning has empirically been observed to occur during the initial phase of training [29, 15]. More concretely, we consider the ERM model ˆ f erm to have weights α erm and β erm, and the regularized model ˆ f reg to have weights α erm and β reg. Second, we assume that weights on the core and spurious features are normalized such that α erm + β erm = α reg + β reg, essentially positing that the proportion of junk feature strength is the same in both ERM and regularized models. This is also justified by the empirical observation that regularization methods such as early stopping make minimal difference to in-distribution accuracy [49].

Main result. The following result compares a minority group test example and a majority group test example with features that are equal in magnitude, and shows that the disagreement is higher for minority group points regardless of whether the ERM or regularized model weights the spurious feature higher. We will use the total variation distance (TVD) as our cost function for disagreement SELF . We focus on TVD instead of KL divergence to simplify the proof; we show in Table 9 (detailed in Appendix B) that the empirical results are competitive using TVD (though worse on CelebA). We show that the difference of the TVD terms is proportional to | β erm -β reg | := | α erm -α reg | , illustrating that the models disagree more drastically the greater the discrepancy between their dependence on the spurious feature (or, equivalently, the core feature).

Theorem 1 (Disagreement SELF ) . Consider two examples x min ∈ X min and x maj ∈ X maj which have the same labels and whose features have the same magnitude, i.e., φ core ( x min ) = φ core ( x maj ) , | φ sp ( x min ) | = | φ sp ( x maj ) | , and ∑ d -2 k =1 φ k junk ( x min ) = ∑ d -2 k =1 φ k junk ( x maj ) . Let P min and Q min be distributions over {-1 , 1 } which take value 1 with probability ( b ( ˆ f erm ( x min ))+1) / 2 and ( b ( ˆ f reg ( x min ))+1) / 2 respectively, and likewise for P maj and Q maj . 7 Then, we have

Proof. For any x ∈ X define P ( x ) = ( b ( ˆ f erm ( x )) + 1) / 2 and Q ( x ) = ( b ( ˆ f reg ( x )) + 1) / 2 . By the assumptions on the features,

Let c = | φ core ( x maj ) | and d = | φ sp ( x maj ) | . By definition of the minority and majority groups,

Together with b ≥ 0 , Equations 7 and 8 show the result.

D Training details

We utilize a ResNet-50 [24] pretrained on ImageNet-1K [57] for the vision tasks and a BERT [14] model pretrained on Book Corpus [83] and English Wikipedia for the language tasks. These pretrained models are used as the initialization for ERM on the four datasets we study. We use standard ImageNet normalization with standard flip and crop data augmentation for the vision tasks and BERT tokenization for the language tasks [28]. Our implementation uses the following packages: NumPy [22], PyTorch [53], Lightning [72], TorchVision [44], Matplotlib [26], Transformers [74], and Milkshake [36].

For ERM and last-layer retraining (Section 4), we do not vary any hyperparameters; their fixed values are listed in Table 11. Recall that our SELF method (Section 5) splits the validation dataset in half, using the first half as a held-out dataset for selecting reweighting points and the second half for model selection [33, 28]. We vary the reweighting dataset size n (the number of points selected from the held-out dataset) in the range (20 , 100 , 500) . For the vision datasets, we search over learning rates (10 -4 , 10 -3 , 10 -2 ) and train for 250 steps. For the language datasets, we search over learning rates (10 -6 , 10 -5 , 10 -4 ) and train for 500 steps. For the two early-stopping methods, we search over checkpoints saved at 10% , 20% , and 50% of model training, while for dropout disagreement, we search over dropout probabilities of 0 . 5 , 0 . 7 , and 0 . 9 applied only at the last layer. See Appendix B for additional ablation studies.

E Additional tables

In this section we provide results in a tabular form to accompany figures in the main paper.

F Broader impacts, limitations, and compute

Broader impacts. We hope our work contributes to the safe and equitable application of machine learning and motivates further research in ML fairness. A potential negative outcome may arise if practitioners assume their models are bias-free after applying our techniques; while we show reduced dependence on spurious correlations, no method can completely alleviate the issue, and other modes of bias may exist. We encourage practitioners to conduct rigorous examinations of model fairness.

Limitations. Our methods take advantage of the specificity of the spurious correlation setting, and therefore would likely underperform on datasets which exhibit a more extreme complete correlation ( i.e., contain zero minority group data) [52, 40]. Furthermore, following previous work in this setting [58, 41, 48, 33, 28], SELF utilizes a small validation dataset with group annotations for model selection. While we show in Appendix B that SELF is robust to using as few as 1% of these annotations, and our last-layer retraining method (Section 4) does not require them, completely removing this assumption is an important direction for future research.

Compute. Our experiments were conducted on Nvidia Tesla V100 and A5000 GPUs. We used about $25000 in compute credits in the course of this research. We believe this paper could be reproduced for under $5000 in credits, and a majority of this compute would go towards the ablation studies. With that said, our last-layer retraining methods only train the linear classifier, making them cheap and efficient to run even on older hardware.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper tackles the aspect of last-layer retraining for better group robustness. The authors show that last-layer retraining can greatly improve worst-group accuracy with little worst-group data. Motivated by this, selective last-layer finetuning (SELF) is proposed to improve group robustness.\n\nSTRENGTHS:\n1. The paper focuses on the interesting aspect of improving robustness via last-layer retraining under spurious correlation. The problem is specific and practical.\n2. The experiments are well done. Also, the comparisons are consistent through out the paper.\n3. The findings and proposed methods are well presented. The finding of \"last-layer retraining can substantially improve WGA even when the reweighting dataset has only a small proportion of worst group data\" is interesting.\n\nWEAKNESSES:\n1. Novelty: \n\n(1) Last-layer retraining/finetuning is a common techqniue used in imbalanced/long-tailed learning and various properties were found, see [1][2]. This put doubt on the novelty of the claimed finding: \"holding out a subset of training data to retrain the last layer can substantially outperform ERM on the entire dataset with no additional data, annotations, or computation for training\". I wonder what are the differences if we treat tail classes in [1][2] as the worst group.\n\n(2) The proposed method: Disagreement-based methods were proposed and they were only adopted to the setting with spurious correlation. This adaptation is fine, but the performance in Table 4 cannot convince me since SELF cannot beat DFR in 2 out of 4 tasks.\n\n2. Writing: \n\n(1) Figure 1(a) misses legends. For Figure 1(b), please note that y-axis means relative increase to avoid confusion. \n\n(2) The contribution part is too long.\n\n[1] Decoupling Representation and Classifier for Long-tailed Recognition. ICLR 2020\n[2] BBN: Bilateral-Branch Network with Cumulative Learning for Long-Tailed Visual Recognition. CVPR 2020\n\nQUESTIONS:\nIs worst-group accuracy a valid metric for evaluating robustness? A common phenomenon in long-tailed learning is that increasing tail-class accuracy often leads to decrease in head class accuracy. I recommend also adding overall accuracy as a metric for comparison. Also, WGA cannot reflect the uniformness of the overall performance.\n\nLIMITATIONS:\nThe limitations are included in the appendix.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes to collect a dataset from misclassifications or disagreements to fine-tune a classifier for improving sub-group accuracy in presence of spurious correlations.\n\nSTRENGTHS:\nThe topic that this paper attempts to address is important. The paper is well-written as well. There are some interesting findings in the paper as well such as finetuning using a dataset collected from disagreements could potentially improve group-accuracy.\n\nWEAKNESSES:\nThe level of novelty presented in the paper is insufficient. My primary concern stems from the method's reliance on group labels during model selection, also known as the validation set. This requirement renders these techniques nearly impractical in real-world scenarios, as defining groups is a non-trivial task even in datasets with just a hundred samples, and this becomes much more challenging for current datasets comprising billions of data points.\n\nThe Waterbirds dataset utilized in this study has been identified to contain certain bugs, which were subsequently addressed in the MaskTune paper [1]. Consequently, the current values provided may be misleading. It is worth noting that the MaskTune paper, contrary to the assumption, does not incorporate any group labels during training or model selection (validation). It would be valuable for the authors to present the worst group accuracies achieved by SELF without the utilization of any training or validation group labels. Additionally, I encourage the authors to compare these results with the methods outlined in Tables 1 and 2 of [1].\n\n[1] Asgari, Saeid, et al. \"Masktune: Mitigating spurious correlations by forcing to explore.\" Advances in Neural Information Processing Systems 35 (2022)\n\nQUESTIONS:\nPlease refer to the Weaknesses section.\n\nLIMITATIONS:\nPlease refer to the Weaknesses section.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper provides a detailed analysis of the DFR procedure [1]. The authors specifically consider the case when the group annotations are limited or not available during training, and show that it is still possible to largely remover the DFR performance. Specifically, the authors retrain the last layer on data with class balancing, and on data where examples are selected according to some rule (e.g. misclassified examples). The authors also provide many detailed ablations and understanding experiments that support their findings.\n\nSTRENGTHS:\n**S1.** The authors provide new interesting insights into the DFR procedure. For example, the results in Figure 1 on the number of minority data needed are quite surprising: even with a small increase in the proportion of minority data it is possible to significantly improve worst-group performance. Also the results in Table 3 on the effectiveness of last-layer retraining without group rebalancing are thought-provoking. Finally, the results in Figure 3, which show that better group balancing does not always lead to better WGA.\n\n**S2.** The authors propose new practical methods (SELF), which build on DFR, but relax the requirement for group-balanced data. The results are discussed in Section 5. The authors also perform detailed ablations on the design decisions for the method, such as what cost function $c$ to use for selecting a subset of the data for last layer retraining.\n\n**S3.** SELF can work even when the reweighting data doesn’t have class labels, and request the labels on a small subset.\n\nOverall, I believe this paper provides new insights that expand our understanding of DFR and more generally training models robust to spurious features. The proposed methodology is also promising.\n\nWEAKNESSES:\nI believe there is one important missing experiment, and several technical details that should be explained or corrected. I explain these in detail below.\n\n**W1.** Experiment: performance vs total number of group labeled examples\n\nAs every other group robustness method, SELF requires group-labeled data to perform hyper-parameter tuning. My understanding is that in all of the experiments, half of the validation data is used for hyper-parameter tuning. In other words, it is not true that the method can work without any group annotations (which is also true of JTT, CnC, and other methods). Moreover, SELF seems to involve more hyper-parameters than DFR (which only tunes regularization strength): length of last layer retraining, size of the reweighting dataset, learning rate, etc.\n\nConsequently, I believe it is important to evaluate performance as a function of the total number of the group-labeled datapoints used by the method. In case of DFR, these datapoints would be used for last-layer retraining, and in case of  SELF, they will only be used for hyper-parameter tuning. I think it may be reasonable to create a group-balanced validation set of size $n$, and plot WGA vs the size $n$. For example of a similar experiment, see Fig. 5 in [2].\n\nI believe, this experiment is quite important for evaluating the methodological contribution of the paper.\n\n**W2.** Civilcomments\n\nI blieve, the results presented in Table 1 for CivilComments mix two different versions of the dataset. Specifically, JTT and CnC use the [Wilds version](https://wilds.stanford.edu/datasets/#civilcomments) with 16 overlapping groups, while DFR, RWY-ES, Group-DRO and the methods described in the paper use a version with 4 groups, where all the spurious attributes are combined together. As a result, there is a big discrepancy in performance. I would recommend either (1) reporting the results for all methods on the WILDS version of the dataset, or (2) removing the JTT and CnC entries and clarifying the difference in the versions of the data.\n\n\n**W3.** Waterbirds\n\nOn waterbirds, the validation data is group balanced up to class balancing. As a result, by group-balancing, the authors are able to achieve DFR performance with class-balanced last-layer retraining (CB in Table 1). However, in this case the method is virtually identical to DFR.\n\nI think it would be more meaningful to do the experiment using a base reweighting dataset that has the same group distribution as the training data. This is the case for all other datasets.\n\nQUESTIONS:\n**Q1. Early stopping:** What is the early stopping criterion for the models in Section 5?\n\n**Q2. Free lunch:** For the free-lunch results in Section 4.2 do you need regularization or early-stopping, optimized on validation WGA? Or is it quite robust to hyper-parameters?\n\n**Q3. Hyper-parameters:** Could you please list all hyper-parameters that a practitioner needs to tune for the SELF method?\n\nFinally, the paper [2] appears to be quite relevant, as it also attempts to automatically construct a reweighting dataset for DFR. However, [2] is a concurrent work, as it was published only after the submission of this paper. So there are no issues with the novelty of the paper.\n\n\n**References**\n\n[1] [_Last Layer Re-Training is Sufficient for Robustness to Spurious Correlations_](https://openreview.net/forum?id=Zb6c8A-Fghk); P. Kirichenko, P. Izmailov, A. G. Wilson; ICLR 2023\n\n[2] [_Simple and Fast Group Robustness by Automatic Feature Reweighting_](https://arxiv.org/abs/2306.11074); S. Qiu, A. Potapczynski, P. Izmailov, A. G. Wilson; ICML 2023\n\nLIMITATIONS:\nLimitations are adequately discussed by the authors.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work tackles an important problem of preventing the reliance of neural networks on spurious correlations. It builds on top of the work [1] primarily by using last layer re-training on class balanced held out dataset without the need for group annotations. They also additionally propose a simple but effective method SELF wherein the samples in disagreement with regularized models in the prediction with class annotations can be used for re-training the last layer. This method is particularly useful when group imbalance is large. \nResults are shown on Waterbirds, CelebA, MultiNLI and CivilComments datasets and performance is comparable to DFR [1].\n\nSTRENGTHS:\n(1) The paper is well written and does very good analysis with ablation studies. \n\n(2) Makes an important observation that it is ok if re-weighting dataset has small proportion of worst group data. This proves that the class balancing is more important than the group balancing. \n\n(3) The paper addresses some of the practical challenges (in terms of computation and training overhead) by eliminating need to annotated group labels which is expensive, need to train twice, need to train the original ERM on class-balanced dataset. \n\n(4)The provided solutions are simple, and will not introduce any implementation or annotation overhead for already trained large scale models. \nThe paper addresses the limitation of accuracy gap when the group imbalance is high and proposes a solution for the same with SELF technique.\n\nWEAKNESSES:\n(1)The observation of class-balancing importance and re-evaluation of group balancing’s role in spurious correlation is a valuable contribution. However, solution in itself lacks novelty as it is not proposing a new algorithm per se. It seems very similar to [1] \"Last layer re-training is sufficient for robustness to spurious correlations\"  with lesser constraints in terms of re-training dataset. \n\n(2) This particular line under section 5  “In addition to the balance of the reweighting, dataset, it is likely that characteristics of the specific data selected also contribute to SELF results. “ is unclear.  Can you provide some qualitative examples from Celeb A where drop-out SELF disagreement does better than misclassification. More insights on what aspects of data makes the dropout SELF better would be good. Similarly, more qualitative example on the “uncertain data points” selected by the SELF-disagreement would provide more valuable insights. A lot of the observations are not supported by reasoning. \n\n(3) Since the key claim is the advantage of not needing group labels annotations - ablation studies on what is the overhead caused for obtaining group annotations (probably annotation time) will help quantify the effectiveness of the proposed method.\nSimilarly ablation studies on the training compute time / mem consumotion between [1] and your method can help quantify the gains if any in training ERM on 95% split and using a MLP instead of Logistic regression.\n\n\nReferences:\n[1] \"Last layer re-training is sufficient for robustness to spurious correlations\"\n\nQUESTIONS:\nCovered in weaknesses.\n\nLIMITATIONS:\nA lot of the questions posed by the paper are still not addressed or analyzed . For e.g., It would be great to see more intuitive and theoretical reasoning on fundamental observations like: why last layer re-training and retraining on held out small class balanced data is improving the group robustness. Though experimental evidence has been provided, further analysis on the theoretical reasoning behind this would make a significant contribution in this subject.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-08-05"}