SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Tyler LaBonte $^{1}$  Vidya Muthukumar $^{2,1}$  Abhishek Kumar $^{3}$
$^{1}$ H. Milton Stewart School of Industrial and Systems Engineering, Georgia Institute of Technology  $^{2}$ School of Electrical and Computer Engineering, Georgia Institute of Technology
3Google DeepMind
{tlabonte, vmuthukumar8}@gatech.edu abhishk@google.com
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.
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).
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. 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 ERM training 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.
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.
<sup>1</sup>Our methods in Section 4 do not require group annotations at all, while our SELF method in Section 5 assumes access to a small validation set with group annotations for model selection following previous work [58, 41, 48, 33, 28]. We show in Appendix B that SELF is robust even with  $1\%$  of these annotations.
Table 1: Comparison to other group robustness methods. We discuss our DFR implementation in Section 3, and our proposed methods of class-balanced (CB) last-layer retraining and early-stop (ES) disagreement SELF in Sections 4 and 5, respectively. Class-balanced ERM is trained on the combined training and held-out datasets. DFR uses group annotations on the held-out dataset, while Group DRO-ES requires them for the training dataset. ES disagreement SELF uses class annotations on the training dataset (for ERM), but requests as few as 20 labels from the held-out dataset. All methods except ERM and CB last-layer retraining use a small set of group annotations for model selection. We list the mean and standard deviation over three independent runs.
Method	Annotations	Worst-group test accuracy
Group	Class	Waterbirds	CelebA	CivilComments	MultiNLI
Class-balanced ERM	X	✓	81.9±3.4	67.2±5.6	61.4±0.7	69.2±1.6
JTT [41, 27]	X	✓	85.6±0.2	75.6±7.7	-	67.5±1.9
RWY-ES [27, 28]	X	✓	74.5±0.0	76.8±7.7	78.9±1.0	68.0±0.4
CnC [81]	X	✓	88.5±0.3	88.8±0.9	-	-
CB last-layer retraining	X	✓	92.6±0.8	73.7±2.8	80.4±0.8	64.7±1.1
ES disagreement SELF	X	X	93.0±0.3	83.9±0.9	79.1±2.1	70.7±2.5
DFR (our impl.)	✓	✓	92.4±0.9	87.0±1.1	81.8±1.6	70.8±0.8
DFR [33, 28]	✓	✓	91.1±0.8	89.4±0.9	78.8±0.5	72.6±0.3
Group DRO-ES [58, 28]	✓	✓	90.7±0.6	90.6±1.6	80.4	73.5
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 retransmits 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 in-distribution 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
Table 2: Dataset composition. We study four well-established benchmarks for group robustness across vision and language tasks. The class probabilities change dramatically when conditioned on the spurious feature. Note that Waterbirds is the only dataset that has a distribution shift and MultiNLI is the only dataset which is class-balanced a priori. Probabilities may not sum to 1 due to rounding.
Dataset	Group g	Training distribution p̂	Data quantity
Class y	Spurious s	p(y)	p(g)	p(y|s)	Train	Val	Test
Waterbirds	landbird	land	.768	.730	.984	3498	467	2225
landbird	water	.038	.148	184	466	2225
waterbird	land	.232	.012	.016	56	133	642
waterbird	water	.220	.852	1057	133	642
CelebA	non-blond	female	.851	.440	.758	71629	8535	9767
non-blond	male	.411	.980	66874	8276	7535
blond	female	.149	.141	.242	22880	2874	2480
blond	male	.009	.020	1387	182	180
CivilComments	neutral	no identity	.887	.551	.921	148186	25159	74780
neutral	identity	.336	.836	90337	14966	43778
toxic	no identity	.113	.047	.079	12731	2111	6455
toxic	identity	.066	.164	17784	2944	8769
MultiNLI	contradiction	no negation	.333	.279	.300	57498	22814	34597
contradiction	negation	.054	.761	11158	4634	6655
entailment	no negation	.334	.327	.352	67376	26949	40496
entailment	negation	.007	.104	1521	613	886
neither	no negation	.333	.323	.348	66630	26655	39930
neither	negation	.010	.136	1992	797	1148
eralization [30], as well as between model classes (e.g., CNNs and Transformers) to predict out-of-distribution 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.
Setting. We consider classification tasks with input domain  $\mathcal{X}$  and target classes  $\mathcal{V}$ . We assume  $S$  is a set of spurious features such that each sample  $x\in \mathcal{X}$  is associated with exactly one feature  $s\in S$ . In conjunction with the target class, the spurious features partition the dataset into groups  $\mathcal{G} = \mathcal{V}\times 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  $\mathcal{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  $\mathcal{G}$  or classes in  $\mathcal{V}$ , 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 \sim \mathrm{Unif}(\mathcal{S})$ , then sample  $x \sim \hat{p}(\cdot | s)$  where  $\hat{p}$  is the training distribution; class balancing is the same with  $\mathcal{V}$  instead of  $\mathcal{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  $\ell_1$ -regularization. The results are averaged over 10 randomly subsampled reweighting datasets, and a hyperparameter search is performed over  $\ell_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  $\ell_2$  regularization instead of searching over  $\ell_1$  regularization strength. We observed similar performance for  $\ell_1$  and  $\ell_2$  regularization, which we believe is because  $\ell_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 \times$  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. 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  $\mathcal{V}$  instead of the groups  $\mathcal{G}$ .
2We note that the Waterbirds dataset is known to contain incorrect labels [66]. We report results on the original, un-corrected version for a fair comparison with previous work.
This version of CivilComments has four groups, used in this work and by [58, 28, 33]. There is another version where the identity categories are not collapsed into one spurious feature; this version is used by [41, 81], so we do not report their CivilComments accuracies in Table 1. Both versions use the WILDS split [35].
4The Waterbirds validation and test datasets still contain more landbirds than waterbirds, but each class has equal quantities of land and water backgrounds.
(a) Worst-group test accuracy
(b) Relative  $\%$  of max WGA increase over ERM
Figure 1: How much worst-group data does last-layer retraining really need? We perform an ablation on the percentage of worst-group data used for class-balanced last-layer retraining, while keeping the total data constant. The results 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, and that class balancing can be a major factor in its performance. The underperformance on Waterbirds below  $20\%$  is because there is too little worst-group data to observe consistent model behavior (less than ten samples). The stars  $\star$  denote the baseline percentage of worst-group data in the training dataset. We plot the mean over three independent runs. See Tables 12 and 13 for details.
Table 3: Last-layer retraining on the held-out dataset. While unbalanced (UB) last-layer retraining decreases performance, class-balanced (CB) last-layer retraining nearly matches DFR on Waterbirds and CivilComments. However, it still trails DFR on CelebA and MultiNLI; we improve these results with the SELF method described in Section 5. CB ERM is trained on the combined training and held-out datasets using class-balanced minibatches. We list the mean and standard deviation over three independent runs.
Method	Group annotations	Worst-group test accuracy
Waterbirds	CelebA	CivilComments	MultiNLI
Class-balanced ERM	X	81.9±3.4	67.2±5.6	61.4±0.7	69.2±1.6
UB last-layer retraining	X	88.0±0.8	41.9±1.4	57.6±4.2	64.6±1.0
CB last-layer retraining	X	92.6±0.8	73.7±2.8	80.4±0.8	64.7±1.1
DFR (our impl.)	✓	92.4±0.9	87.0±1.1	81.8±1.6	70.8±0.8
DFR [33, 28]	✓	91.1±0.8	89.4±0.9	78.8±0.5	72.6±0.3
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.