On the Unreasonable Effectiveness of Last-layer Retraining

TL;DR

Last-layer retraining (LLR) improves neural network robustness and minority group performance, primarily due to better group balance in the held-out set rather than neural collapse mitigation.

Contribution

The paper challenges the neural collapse hypothesis and provides evidence that group balance in the held-out set explains LLR's effectiveness, highlighting implicit group-balancing algorithms.

Findings

LLR improves worst-group accuracy even with imbalanced held-out sets

Neural collapse does not explain LLR's effectiveness

Implicit group balancing in algorithms like CB-LLR and AFR enhances robustness

Abstract

Last-layer retraining (LLR) methods -- wherein the last layer of a neural network is reinitialized and retrained on a held-out set following ERM training -- have garnered interest as an efficient approach to rectify dependence on spurious correlations and improve performance on minority groups. Surprisingly, LLR has been found to improve worst-group accuracy even when the held-out set is an imbalanced subset of the training set. We initially hypothesize that this ``unreasonable effectiveness'' of LLR is explained by its ability to mitigate neural collapse through the held-out set, resulting in the implicit bias of gradient descent benefiting robustness. Our empirical investigation does not support this hypothesis. Instead, we present strong evidence for an alternative hypothesis: that the success of LLR is primarily due to better group balance in the held-out set. We conclude by showing how the recent algorithms CB-LLR and AFR perform implicit group-balancing to elicit a robustness improvement.