# Mind the GAP: Improving Robustness to Subpopulation Shifts with Group-Aware Priors

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 23 of its 41 references that are in the map.
- First posted on arXiv 14 Mar 2024.

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

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Decoupled Weight Decay Regularization](https://mlsanity.com/paper/1711.05101/) (2017)
- [Sharpness-Aware Minimization for Efficiently Improving Generalization](https://mlsanity.com/paper/2010.01412/) (2020)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [Just Train Twice: Improving Group Robustness without Training Group Information](https://mlsanity.com/paper/2107.09044/) (2021)
- [Bayesian Deep Learning and a Probabilistic Perspective of Generalization](https://mlsanity.com/paper/2002.08791/) (2020)
- [Last Layer Re-Training is Sufficient for Robustness to Spurious Correlations](https://mlsanity.com/paper/2204.02937/) (2022)
- [Simple data balancing achieves competitive worst-group-accuracy](https://mlsanity.com/paper/2110.14503/) (2021)
- [Correct-N-Contrast: A Contrastive Approach for Improving Robustness to Spurious Correlations](https://mlsanity.com/paper/2203.01517/) (2022)
- [On Feature Learning in the Presence of Spurious Correlations](https://mlsanity.com/paper/2210.11369/) (2022)
- [Distributionally Robust Language Modeling](https://mlsanity.com/paper/1909.02060/) (2019)
- [Change is Hard: A Closer Look at Subpopulation Shift](https://mlsanity.com/paper/2302.12254/) (2023)
- [Spread Spurious Attribute: Improving Worst-group Accuracy with Spurious Attribute Estimation](https://mlsanity.com/paper/2204.02070/) (2022)
- [Coping with Label Shift via Distributionally Robust Optimisation](https://mlsanity.com/paper/2010.12230/) (2020)
- [Simple and Fast Group Robustness by Automatic Feature Reweighting](https://mlsanity.com/paper/2306.11074/) (2023)
- [Pre-Train Your Loss: Easy Bayesian Transfer Learning with Informative Priors](https://mlsanity.com/paper/2205.10279/) (2022)
- [Does Distributionally Robust Supervised Learning Give Robust Classifiers?](https://mlsanity.com/paper/1611.02041/) (2016)
- [Discovering environments with XRM](https://mlsanity.com/paper/2309.16748/) (2023)
- [BARACK: Partially Supervised Group Robustness With Guarantees](https://mlsanity.com/paper/2201.00072/) (2021)
- [Function-Space Regularization in Neural Networks: A Probabilistic Perspective](https://mlsanity.com/paper/2312.17162/) (2023)
- [Is Last Layer Re-Training Truly Sufficient for Robustness to Spurious Correlations?](https://mlsanity.com/paper/2308.00473/) (2023)

## Abstract (arXiv)

Machine learning models often perform poorly under subpopulation shifts in the data distribution. Developing methods that allow machine learning models to better generalize to such shifts is crucial for safe deployment in real-world settings. In this paper, we develop a family of group-aware prior (GAP) distributions over neural network parameters that explicitly favor models that generalize well under subpopulation shifts. We design a simple group-aware prior that only requires access to a small set of data with group information and demonstrate that training with this prior yields state-of-the-art performance -- even when only retraining the final layer of a previously trained non-robust model. Group aware-priors are conceptually simple, complementary to existing approaches, such as attribute pseudo labeling and data reweighting, and open up promising new avenues for harnessing Bayesian inference to enable robustness to subpopulation shifts.