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

Tim G. J. Rudner, Ya Shi Zhang, Andrew Gordon Wilson, Julia Kempe
Proceedings of The 27th International Conference on Artificial Intelligence and Statistics, PMLR 238:127-135, 2024.

Abstract

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v238-rudner24a, title = {Mind the {GAP}: Improving Robustness to Subpopulation Shifts with Group-Aware Priors}, author = {Rudner, Tim G. J. and Shi Zhang, Ya and Wilson, Andrew Gordon and Kempe, Julia}, booktitle = {Proceedings of The 27th International Conference on Artificial Intelligence and Statistics}, pages = {127--135}, year = {2024}, editor = {Dasgupta, Sanjoy and Mandt, Stephan and Li, Yingzhen}, volume = {238}, series = {Proceedings of Machine Learning Research}, month = {02--04 May}, publisher = {PMLR}, pdf = {https://proceedings.mlr.press/v238/rudner24a/rudner24a.pdf}, url = {https://proceedings.mlr.press/v238/rudner24a.html}, abstract = {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.} }
Endnote
%0 Conference Paper %T Mind the GAP: Improving Robustness to Subpopulation Shifts with Group-Aware Priors %A Tim G. J. Rudner %A Ya Shi Zhang %A Andrew Gordon Wilson %A Julia Kempe %B Proceedings of The 27th International Conference on Artificial Intelligence and Statistics %C Proceedings of Machine Learning Research %D 2024 %E Sanjoy Dasgupta %E Stephan Mandt %E Yingzhen Li %F pmlr-v238-rudner24a %I PMLR %P 127--135 %U https://proceedings.mlr.press/v238/rudner24a.html %V 238 %X 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.
APA
Rudner, T.G.J., Shi Zhang, Y., Wilson, A.G. & Kempe, J.. (2024). Mind the GAP: Improving Robustness to Subpopulation Shifts with Group-Aware Priors. Proceedings of The 27th International Conference on Artificial Intelligence and Statistics, in Proceedings of Machine Learning Research 238:127-135 Available from https://proceedings.mlr.press/v238/rudner24a.html.

Related Material