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.

翻译：机器学习模型在数据分布发生子群体偏移时往往表现不佳。开发能使其在此类偏移下更好地泛化的方法，对于在现实场景中安全部署至关重要。本文提出了一类基于神经网络参数的组感知先验（GAP）分布，其明确偏向于在子群体偏移下具有良好泛化能力的模型。我们设计了一种简单的组感知先验，仅需访问少量带有组信息的数据，并证明使用该先验进行训练即可达到最先进的性能——即便仅对预先训练过的非鲁棒模型的最后一层进行重训练。组感知先验概念简洁，可与现有方法（如属性伪标签、数据重加权）互补，为利用贝叶斯推断实现子群体偏移鲁棒性开辟了有前景的新途径。