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

Tim G.J. Rudner, Ya Shi Zhang, Andrew Wilson, Julia Kempe

Research output: Contribution to journalConference articlepeer-review


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.

Original languageEnglish (US)
Pages (from-to)127-135
Number of pages9
JournalProceedings of Machine Learning Research
StatePublished - 2024
Event27th International Conference on Artificial Intelligence and Statistics, AISTATS 2024 - Valencia, Spain
Duration: May 2 2024May 4 2024

ASJC Scopus subject areas

  • Artificial Intelligence
  • Software
  • Control and Systems Engineering
  • Statistics and Probability


Dive into the research topics of 'Mind the GAP: Improving Robustness to Subpopulation Shifts with Group-Aware Priors'. Together they form a unique fingerprint.

Cite this