SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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. 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.
This is necessary because we keep the total data constant. By reducing the quantity of worst-group data to zero, we will at most double the size of any other group.
(a) Training dataset only
(b) Combined training and held-out datasets
Figure 2: A "free lunch" in group robustness. We compare class-balanced (CB) ERM on the entire dataset to splitting the dataset and performing CB ERM on the first  $(95\%)$  split and CB last-layer retraining on the second  $(5\%)$  split. This technique improves worst-group accuracy on Waterbirds, CelebA, and CivilComments by up to  $17\%$  while using no additional data or annotations for training beyond ERM. We believe it underperforms on MultiNLI because there is not enough data in the first split, i.e., ERM performance can be improved by collecting more data. We plot the mean and standard deviation over three independent runs. See Table 14 for detailed results.
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  $\star$  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
Table 4: Comparison of selective last-layer finetuning methods. SELF nearly matches DFR and improves WGA over class-balanced (CB) last-layer retraining by up to  $12\%$ . Early-stop (ES) disagreement SELF performs the best overall, and disagreement methods perform especially well on CivilComments. Dropout and ES disagreement request few as 20 class annotations from the held-out dataset. All SELF methods use a small set of group annotations for model selection. We list the mean and standard deviation over three independent runs.
Method	Held-out 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
CB last-layer retraining	X	✓	92.6±0.8	73.7±2.8	80.4±0.8	64.7±1.1
Random SELF	X	X	91.1±1.9	80.2±6.4	80.5±1.6	65.0±4.0
Misclassification SELF	X	✓	92.6±0.8	83.0±6.1	62.7±4.6	72.2±2.2
ES misclassification SELF	X	✓	92.2±0.7	80.4±3.9	65.8±7.6	73.3±1.2
Dropout disagreement SELF	X	X	92.3±0.5	85.7±1.6	69.9±5.2	68.7±3.4
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
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 \pm 2.4\%$  WGA on the training dataset and  $69.2 \pm 1.6\%$  on the combined training and held-out datasets, while our DFR implementation achieves  $70.8 \pm 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.
A critical 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: \mathcal{X} \to \mathbb{R}^{|\mathcal{Y}|}$  be models with logit outputs,  $X \subseteq \mathcal{X}$  be a held-out dataset,  $c: \mathbb{R}^{2|\mathcal{Y}|} \to \mathbb{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.,
$$
D = \underset {S \subseteq X, | S | = n} {\operatorname {a r g m a x}} \sum_ {x \in S} c (f (x), g (x)). \tag {1}
$$
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.
<sup>6</sup>Using early-stop misclassification to upsample worst-group data is the premise of JTT [41].
Figure 3: SELF upsamples the worst group. We plot the percentage of the reweighting dataset consisting of worst-group data for each SELF method. "Baseline" represents the percentage of worst-group data in the held-out dataset, while "Balanced" is the percentage of worst-group data necessary to achieve group balance. This is to the best of our knowledge the first empirical evidence that model disagreement is an effective method for upsampling worst-group data. The worst groups in each dataset are listed in Section 4.1. We plot the mean for the best dataset size  $n$  over three independent runs. See Table 15 for detailed results.
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.
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.
[82] 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.
[83] 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.
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 includes 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.
Table 5: Empirical risk minimization. First, we compare class-unbalanced (CU) and class-balanced (CB) ERM on the training dataset vs. the combined training and held-out datasets (i.e., the training dataset plus half the validation dataset). We list the mean and standard deviation over three independent runs. Note that the Waterbirds validation dataset has a different distribution than the training dataset (and, in particular, is group-balanced), and that MultiNLI is class-balanced a priori.
Method	Held-out dataset included	Worst-group test accuracy
Waterbirds	CelebA	CivilComments	MultiNLI
CU ERM	X	72.4±1.0	44.1±0.9	63.8±6.2	67.4±2.4
CB ERM	X	72.6±3.2	66.3±3.2	60.2±2.7	67.4±2.4
CU ERM	✓	81.6±1.5	44.5±3.4	59.1±2.2	69.1±1.3
CB ERM	✓	81.9±3.4	67.2±5.6	61.4±0.7	69.2±1.6
Table 6: Balancing methodology comparison. Next, we compare last-layer retraining on the held-out set with different balancing methodologies, each initialized with class-balanced ERM features. In "sampling", we sample from the held-out dataset at a non-uniform rate so that each minibatch is class- or group-balanced in expectation, while in "subset", we train on a random class- or group-balanced subset of the held-out dataset. Specifically, for group-balanced sampling we first sample  $s \sim$  Unif(S), then sample  $x \sim  \widehat{p}\left( {\cdot  \mid  s}\right)$  where  $\widehat{p}$  is the training distribution; class-balanced sampling is the same with  $\mathcal{V}$  instead of  $S$  . For subset balancing,we keep all data from the smallest group/class and downsample the others uniformly at random to that size. We list the mean and standard deviation over three independent runs.
Last-layer retraining method	Worst-group test accuracy
Waterbirds	CelebA	CivilComments	MultiNLI
Class-unbalanced	88.0±0.8	41.9±1.4	57.6±4.2	64.6±1.0
Class-balanced sampling	92.6±0.8	73.7±2.8	80.4±0.8	64.7±1.1
Class-balanced subset	92.1±0.9	74.6±2.0	80.2±1.4	64.5±1.3
Group-balanced sampling	92.4±0.9	87.0±1.1	81.8±1.6	70.8±0.8
Group-balanced subset	91.6±1.9	88.1±1.0	79.2±1.8	68.3±1.8
DFR [33, 28]	91.1±0.8	89.4±0.9	78.8±0.5	72.6±0.3
Table 7: Necessity of the held-out dataset. Next, we investigate whether holding out a subset of the training dataset for class-balanced (CB) last-layer retraining is essential, or if retraining on the entire (previously seen) dataset is also effective. For retraining on the training dataset, we use the same class-balanced last-layer retraining procedure as the held-out dataset, but we train for 20 epochs for the vision tasks and 2 epochs for the language tasks. For retraining on the held-out dataset, we report the best over four splits (i.e., the same numbers as Figure 2). Our results suggest that holding out data is necessary to achieve maximal worst-group accuracy, though last-layer retraining on the training dataset interestingly prevents the performance decrease on MultiNLI. We list the mean and standard deviation over three independent runs.
Method	Worst-group test accuracy
Waterbirds	CelebA	CivilComments	MultiNLI
CB ERM	72.6±3.2	66.3±3.2	60.2±2.7	67.4±2.4
CB last-layer retraining on training dataset	71.0±2.5	66.9±1.4	61.9±0.8	67.0±1.5
CB last-layer retraining on held-out dataset	77.4±0.3	73.0±2.3	77.9±1.5	63.0±1.5
Table 8: Average accuracy performance. We detail the average test accuracy of our methods on the 4 benchmark datasets. Both of our methods have similar average accuracy to DFR, which experiences a slight accuracy/robustness tradeoff compared to ERM (as is typical in the robustness literature). We list the mean and standard deviation over three independent seeds.
Method	Group Anns	Average test accuracy
Waterbirds	CelebA	CivilComments	MultiNLI
ERM	X	90.2±0.7	94.4±0.2	92.0±0.2	81.8±0.2
CB last-layer retraining	X	94.8±0.3	93.6±0.2	87.1±0.0	82.0±0.2
ES disagreement SELF	X	94.0±1.7	91.7±0.4	87.7±0.6	81.2±0.7
DFR (our impl.)	✓	94.9±0.3	92.6±0.5	87.5±0.2	81.7±0.2
Figure 4: Class-balanced last-layer retraining renders class-balanced ERM optional. We compare class-balanced (CB) last-layer retraining on the held-out dataset initialized with unbalanced (UB) or CB ERM features. Our results show that CB ERM is unnecessary as long as the last layer is retrained with class balancing. We plot the mean and standard deviation over three independent runs.
B Additional experiments for Section 5
In this section we include the detailed results of additional experiments for Section 5.
(b) CelebA
(a) Waterbirds
(c) CivilComments