Towards Last-layer Retraining for Group Robustness with Fewer Annotations (2309.08534v3)

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 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.

• The paper shows that last-layer retraining on a class-balanced dataset subset can significantly improve neural network group robustness, achieving near state-of-the-art performance with less worst-group data than required by methods like DFR.
• It introduces a "free lunch" concept where improving group robustness is possible by randomly splitting and balancing existing training data for last-layer retraining, avoiding the need for additional annotations.
• The study proposes Selective Last-layer Finetuning (SELF), which uses model disagreement to identify challenging samples and achieves group robustness comparable to DFR with drastically reduced (less than 3%) annotation requirements.

An Examination of Last-layer Retraining for Improving Group Robustness in Neural Networks

The paper "Towards Last-layer Retraining for Group Robustness with Fewer Annotations" highlights a significant investigation into enhancing the group robustness of neural networks, particularly in the context of spurious correlations that result in underperformance on minority groups. This discussion is crucial for understanding the nuanced dynamics between empirical risk minimization (ERM) and its tendencies to overfit spurious correlations, thereby emphasizing an alternative strategy: last-layer retraining.

ERM, while effective in general scenarios, often leads to models that depend inappropriately on spurious correlations present in training data, hence performing poorly on underrepresented minority groups. This paper explores last-layer retraining as a paradigm to address these deficiencies without the prerequisite of extensive group annotations. The research capitalizes on the approach exemplified by Deep Feature Reweighting (DFR), which achieves state-of-the-art group robustness by retraining the neural network’s last layer on a group-balanced dataset. However, DFR's reliance on complete group annotation limits its applicability in many real-world scenarios.

The authors first conduct an ablation paper to interrogate how much worst-group data is necessary to enhance model performance via retraining. Results demonstrate that significant improvements in worst-group accuracy can be achieved without substantial proportions of worst-group data when the last-layer is class-balanced during retraining. Remarkably, retraining the last layer on a class-balanced dataset garnered 94% of the maximum performance offered by DFR, surpassing traditional ERM configurations by a substantial margin.

In a striking revelation, the paper introduces the idea of a 'free lunch' in group robustness. It posits that improved group accuracy can be attained by randomly splitting available training data and balancing this subset for last-layer retraining, rather than relying on ERM for the entire dataset. This technique intriguingly achieves enhanced performance without the need for additional annotations or data, representing a computational and resource-efficient paradigm shift for practitioners.

Yet, the paper does not conclude with class-balanced approaches. It innovatively proposes Selective Last-layer Finetuning (SELF), a method that eschews both group annotations and large retraining datasets. SELF selects its retraining points based on model disagreement or misclassification, and uniquely employs only the challenging samples for last-layer finetuning. Empirically, SELF nearly matches the performance of DFR while drastically reducing the annotation burden—down to less than 3% of class annotations.

In dissecting the effectiveness of SELF, the work highlights the critical insight that model disagreements are crucial in identifying worst-group samples. Across robust theoretical and empirical analyses, it is shown that disagreements, particularly between early-stopped and ERM models, reliably indicate samples where the model prediction may be inaccurate due to its spurious dependencies. The theoretical underpinning suggests that uncertainty or disagreement consistently upsamples minority group data, thus aligning with the enhanced robustness observed.

While this paper does not claim groundbreaking innovation, its contributions have profound implications on practical AI deployment, particularly in handling bias within datasets without significant annotation overhead. This exploration of retraining strategies presents an avenue for more equitable deployment in fields where data annotations are scarce or costly, such as facial recognition and healthcare AI systems. Future research may further extend these findings, exploring larger LLMs or deeper architectures where last-layer refinements could continue to advance fairness and equity in model outputs.