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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Rudner, Tim G. J., Zhang, Ya Shi, Wilson, Andrew Gordon, Kempe, Julia
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866910368080592896
author Rudner, Tim G. J.
Zhang, Ya Shi
Wilson, Andrew Gordon
Kempe, Julia
author_facet Rudner, Tim G. J.
Zhang, Ya Shi
Wilson, Andrew Gordon
Kempe, Julia
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2403_09869
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Mind the GAP: Improving Robustness to Subpopulation Shifts with Group-Aware Priors
Rudner, Tim G. J.
Zhang, Ya Shi
Wilson, Andrew Gordon
Kempe, Julia
Machine Learning
Artificial Intelligence
Methodology
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.
title Mind the GAP: Improving Robustness to Subpopulation Shifts with Group-Aware Priors
topic Machine Learning
Artificial Intelligence
Methodology
url https://arxiv.org/abs/2403.09869