# Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift

arXiv: https://arxiv.org/abs/2505.23027 · HTML: https://mlsanity.com/paper/2505.23027/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 40 of its 52 references that are in the map.
- First posted on arXiv 29 May 2025.

## Builds on (most cited first, 30 of 40)

- [LAION-5B: An open large-scale dataset for training next generation image-text models](https://mlsanity.com/paper/2210.08402/) (2022)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [Prototypical Networks for Few-shot Learning](https://mlsanity.com/paper/1703.05175/) (2017)
- [Shortcut Learning in Deep Neural Networks](https://mlsanity.com/paper/2004.07780/) (2020)
- [CheXpert: A Large Chest Radiograph Dataset with Uncertainty Labels and Expert Comparison](https://mlsanity.com/paper/1901.07031/) (2019)
- [BDD100K: A Diverse Driving Dataset for Heterogeneous Multitask Learning](https://mlsanity.com/paper/1805.04687/) (2018)
- [Decoupling Representation and Classifier for Long-Tailed Recognition](https://mlsanity.com/paper/1910.09217/) (2019)
- [Scaling Vision with Sparse Mixture of Experts](https://mlsanity.com/paper/2106.05974/) (2021)
- [Mixture-of-Experts with Expert Choice Routing](https://mlsanity.com/paper/2202.09368/) (2022)
- [Just Train Twice: Improving Group Robustness without Training Group Information](https://mlsanity.com/paper/2107.09044/) (2021)
- [Deep Ensembles: A Loss Landscape Perspective](https://mlsanity.com/paper/1912.02757/) (2019)
- [Nuanced Metrics for Measuring Unintended Bias with Real Data for Text Classification](https://mlsanity.com/paper/1903.04561/) (2019)
- [Last Layer Re-Training is Sufficient for Robustness to Spurious Correlations](https://mlsanity.com/paper/2204.02937/) (2022)
- [Confounding variables can degrade generalization performance of radiological deep learning models](https://mlsanity.com/paper/1807.00431/) (2018)
- [Noise or Signal: The Role of Image Backgrounds in Object Recognition](https://mlsanity.com/paper/2006.09994/) (2020)
- [An Investigation of Why Overparameterization Exacerbates Spurious Correlations](https://mlsanity.com/paper/2005.04345/) (2020)
- [Failing Loudly: An Empirical Study of Methods for Detecting Dataset Shift](https://mlsanity.com/paper/1810.11953/) (2018)
- [Hidden Stratification Causes Clinically Meaningful Failures in Machine Learning for Medical Imaging](https://mlsanity.com/paper/1909.12475/) (2019)
- [Simple data balancing achieves competitive worst-group-accuracy](https://mlsanity.com/paper/2110.14503/) (2021)
- [Correct-N-Contrast: A Contrastive Approach for Improving Robustness to Spurious Correlations](https://mlsanity.com/paper/2203.01517/) (2022)
- [On Feature Learning in the Presence of Spurious Correlations](https://mlsanity.com/paper/2210.11369/) (2022)
- [Robust Classification with Convolutional Prototype Learning](https://mlsanity.com/paper/1805.03438/) (2018)
- [Change is Hard: A Closer Look at Subpopulation Shift](https://mlsanity.com/paper/2302.12254/) (2023)
- [NICO++: Towards Better Benchmarking for Domain Generalization](https://mlsanity.com/paper/2204.08040/) (2022)
- [MetaShift: A Dataset of Datasets for Evaluating Contextual Distribution Shifts and Training Conflicts](https://mlsanity.com/paper/2202.06523/) (2022)
- [Agree to Disagree: Diversity through Disagreement for Better Transferability](https://mlsanity.com/paper/2202.04414/) (2022)
- [Simple and Fast Group Robustness by Automatic Feature Reweighting](https://mlsanity.com/paper/2306.11074/) (2023)
- [Diversify and Disambiguate: Learning From Underspecified Data](https://mlsanity.com/paper/2202.03418/) (2022)
- [Towards Last-layer Retraining for Group Robustness with Fewer Annotations](https://mlsanity.com/paper/2309.08534/) (2023)
- [Online Prototype Learning for Online Continual Learning](https://mlsanity.com/paper/2308.00301/) (2023)

## Built on by (most cited first, 5 of 5)

- [Rank-Constrained Adaptation for Reliable Real-World Performance](https://mlsanity.com/paper/2602.06924/) (2026)
- [Beyond Metadata: CAPRA for Hidden Subgroup Analysis under Missing Metadata in Medical Imaging](https://mlsanity.com/paper/2607.09102/) (2026)
- [ProME: Prototype-Margin Environments with Repair-Aware Selection for Group-Robust Learning](https://mlsanity.com/paper/2608.13190/) (2026)
- [From the Loss Landscape to Diverse Feature Learning in Neural Networks](https://mlsanity.com/paper/2608.28948/) (2026)
- [AIA2: Attribute-Agnostic Imbalance Augmentation for Subgroup Robustness](https://mlsanity.com/paper/2608.30297/) (2026)

## Abstract (arXiv)

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