Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift

Motivation

In real-world datasets, classifiers often learn to rely on shortcuts or spurious features that correlate with the labels in the training distribution but fail to hold in deployment. These failures are especially severe when certain subgroups are underrepresented or shift at test time, making it difficult for the model to learn a single decision boundary that generalizes well.

A natural solution is to move beyond a single global decision rule. When a model is encouraged to consider multiple distinct ways to separate the data, it becomes more likely that at least one decision rule will align with the correct semantics across diverse subpopulations. This forms the key insight behind our approach.

Rather than relying on prior knowledge of subgroup structure, we explore the feature space. By training multiple classifiers that are explicitly encouraged to be different from one another, we can uncover a richer set of decision boundaries that are more robust to distribution shifts.

(a) A single model trained on mixed subgroups learns a biased decision boundary favoring the majority. (b–d) As ensemble size increases, each model captures different subgroup patterns, improving classification across subpopulations.

Diversified Prototypical Ensembles

DPE is a two-stage approach for improving robustness to subpopulation shifts by training an ensemble of diverse prototype classifiers. DPE first trains a feature extractor \( f: \mathbb{R}^n \to \mathbb{R}^d \) using standard empirical risk minimization (ERM) on the full training data. This backbone is frozen and then used to extract features for a small, class-balanced subset of the validation set. On top of these features, a diversified ensemble of prototype classifiers is trained. Notably, no subgroup annotations are required in this second stage, though they can improve balance via stratified sampling.

Evaluation Metrics

To assess robustness under subpopulation shift, we primarily use Worst-Group Accuracy (WGA), the standard evaluation metric in this field. WGA quantifies performance on the most challenging or underrepresented subgroup in the dataset, reflecting how well a model generalizes beyond dominant patterns. In addition to WGA, we report Balanced Accuracy (BA) and Standard Accuracy to provide a more complete picture of model performance across all classes.

Generalization Under Distribution Shifts

We evaluate the robustness of DPE against subpopulation shifts using nine real-world datasets across vision, language, and medical domains. These datasets are chosen to capture common challenges such as spurious correlations, attribute imbalance, class imbalance, and attribute generalization, all of which can degrade model performance under distribution shift. The evaluation includes Waterbirds, CelebA, MetaShift, ImageNetBG, NICO++, Living17, CheXpert, CivilComments, and MultiNLI, using the training, validation, and test splits provided by Yang et al. (2023). These datasets help assess model robustness when attribute or subgroup information is available or missing.

Across a wide range of subpopulation shift scenarios, DPE consistently achieves the highest worst-group accuracy, outperforming both classical and state-of-the-art methods. Its effectiveness holds under missing subgroup information, challenging shift types like attribute imbalance and generalization, and across different backbone strengths.

In our experiments, ERM^* refers to a stronger feature extractor trained with standard empirical risk minimization but using enhanced configurations such as longer training schedules, stronger data augmentations, potentially improving the quality of the extracted features.

Worst-group improvement over ERM* when using DPE with and without subgroup annotations.

Ensemble Diversification Strategies

The figure below illustrates how different ensemble diversification strategies impact both worst-group accuracy and balanced accuracy as the number of ensemble members increases. The comparison includes three methods:

1. Training each prototype on a fixed data subset (no diversification)
2. Training on random subsets for each ensemble member
3. Combining random subset selection with inter-prototype similarity loss

The combined strategy consistently yields the best worst-group accuracy across datasets. This highlights the importance of combining both explicit (loss-based) and implicit (sampling-based) diversification to ensure broad coverage of the data distribution and robustness to subpopulation shifts.

Effect of different ensemble diversification methods on performance with increasing ensemble size.

Prototype–Subgroup Alignment

To better understand the semantic structure captured by the Diversified Prototypical Ensemble (DPE), we conducted an exploratory analysis on the Waterbirds dataset. For each prototype, we retrieved its top-10 closest validation samples and analyzed the emerging patterns using ChatGPT. As visualized below, the learned prototypes show consistent alignment with ecologically or visually meaningful subpopulations, even though no subgroup labels were used during training.

In the Waterbirds class, the prototypes capture structured concepts such as aquatic divers, large-bodied birds in flight, and compact seabirds in bamboo-heavy or terrestrial environments. In the Landbirds class, the clusters reflect postural cues, background settings, and spurious correlations like songbirds appearing in human-made or beach scenes.

These findings suggest that DPE facilitates implicit subgroup discovery through diversification, which may contribute to its strong performance on worst-group accuracy.

Each row shows the top-10 validation samples nearest to one of the learned waterbird prototypes. Patterns include aquatic divers, birds in flight, terrestrial seabirds, and upright gulls, as interpreted by ChatGPT.

Each row shows the top-10 validation samples closest to a DPE prototype for the Landbirds class. Learned clusters correspond to postural variation, ecological context, and background-driven biases such as birds in human-made settings.