Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift

2025-05-29

AI Summary

Subgroup distribution shifts severely degrade model performance on the worst-case subgroups, yet existing methods rely on subgroup labels or strong prior assumptions—rendering them impractical in real-world settings. To address this, we propose an unsupervised subgroup generalization framework: leveraging a pretrained feature extractor, we construct a diversity-driven ensemble of prototype classifiers and jointly optimize at both feature and instance levels to adaptively capture heterogeneous subgroup structures. Crucially, our method requires no subgroup annotations, makes no assumptions about the number or semantic meaning of subgroups, and thus achieves, for the first time, fully group-supervision-free robust subgroup generalization. Evaluated across nine diverse real-world datasets spanning multiple domains, our approach achieves substantial average improvements in worst-subgroup accuracy and outperforms state-of-the-art methods on multiple tasks.

Technology Category

• Machine Learning: Ensemble Methods
• Computer Vision: Adversarial Attacks & Robustness
• Natural Language Processing: Safety and Robustness

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs

Abstract

The subpopulationtion shift, characterized by a disparity in subpopulation distributibetween theween the training and target datasets, can significantly degrade the performance of machine learning models. Current solutions to subpopulation shift involve modifying empirical risk minimization with re-weighting strategies to improve generalization. This strategy relies on assumptions about the number and nature of subpopulations and annotations on group membership, which are unavailable for many real-world datasets. Instead, we propose using an ensemble of diverse classifiers to adaptively capture risk associated with subpopulations. Given a feature extractor network, we replace its standard linear classification layer with a mixture of prototypical classifiers, where each member is trained to classify the data while focusing on different features and samples from other members. In empirical evaluation on nine real-world datasets, covering diverse domains and kinds of subpopulation shift, our method of Diverse Prototypical Ensembles (DPEs) often outperforms the prior state-of-the-art in worst-group accuracy. The code is available at https://github.com/minhto2802/dpe4subpop

Problem

• Addressing performance degradation due to subpopulation shift in datasets
• Overcoming limitations of current re-weighting strategies for generalization
• Improving worst-group accuracy with diverse prototypical classifier ensembles

Innovation

• Ensemble of diverse prototypical classifiers
• Replaces linear layer with mixture
• Improves worst-group accuracy