# Theoretical Guarantees of Data Augmented Last Layer Retraining Methods

arXiv: https://arxiv.org/abs/2405.05934 · HTML: https://mlsanity.com/paper/2405.05934/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 10 of its 14 references that are in the map.
- First posted on arXiv 9 May 2024.

## Builds on (most cited first, 10 of 10)

- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Just Train Twice: Improving Group Robustness without Training Group Information](https://mlsanity.com/paper/2107.09044/) (2021)
- [Last Layer Re-Training is Sufficient for Robustness to Spurious Correlations](https://mlsanity.com/paper/2204.02937/) (2022)
- [Improving Out-of-Distribution Robustness via Selective Augmentation](https://mlsanity.com/paper/2201.00299/) (2022)
- [Change is Hard: A Closer Look at Subpopulation Shift](https://mlsanity.com/paper/2302.12254/) (2023)
- [Simple and Fast Group Robustness by Automatic Feature Reweighting](https://mlsanity.com/paper/2306.11074/) (2023)
- [Towards Last-layer Retraining for Group Robustness with Fewer Annotations](https://mlsanity.com/paper/2309.08534/) (2023)
- [Why does Throwing Away Data Improve Worst-Group Error?](https://mlsanity.com/paper/2205.11672/) (2022)
- [Finite sample performance of linear least squares estimation](https://mlsanity.com/paper/1810.06380/) (2018)

## Built on by (most cited first, 3 of 3)

- [Why Loss Re-weighting Works If You Stop Early: Training Dynamics of Unconstrained Features](https://mlsanity.com/paper/2601.12011/) (2026)
- [Thumb on the Scale: Optimal Loss Weighting in Last Layer Retraining](https://mlsanity.com/paper/2506.20025/) (2025)
- [Label Noise Robustness for Domain-Agnostic Fair Corrections via Nearest Neighbors Label Spreading](https://mlsanity.com/paper/2406.09561/) (2024)

## Abstract (arXiv)

Ensuring fair predictions across many distinct subpopulations in the training data can be prohibitive for large models. Recently, simple linear last layer retraining strategies, in combination with data augmentation methods such as upweighting, downsampling and mixup, have been shown to achieve state-of-the-art performance for worst-group accuracy, which quantifies accuracy for the least prevalent subpopulation. For linear last layer retraining and the abovementioned augmentations, we present the optimal worst-group accuracy when modeling the distribution of the latent representations (input to the last layer) as Gaussian for each subpopulation. We evaluate and verify our results for both synthetic and large publicly available datasets.