Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift
Authors: Nguyen Nhat Minh To
Paul F R Wilson
Viet Nguyen
Mohamed Harmanani
Michael Cooper
Fahimeh Fooladgar
Purang Abolmaesumi
Parvin Mousavi
Rahul Krishnan
Venue: ICML 2025 poster (2025)
Keywords: distribution shift
subpopulation shift
spurious correlation
class imbalance
attribute imbalance
Primary area: deep_learning->robustness
TL;DR: This paper proposes a prototype-based ensemble method to improve robustness against subpopulation shifts without requiring subgroup labels, outperforming state-of-the-art methods across several benchmark datasets.

Abstract: Subpopulation shift, characterized by a disparity in subpopulation distribution between 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.

Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift
Minh Nguyen Nhat To Paul F R Wilson Viet Nguyen 4 Mohamed Harmanani* Michael Cooper
Fahimeh Fooladgar' Purang Abolmaesumi
Parvin Mousavi 32 Rahul G. Krishnan 42
Abstract
tween training and target distributions. The study quantita- tively categorized subpopulation shifts into four fundamen- tal types: (1) Spurious Correlations non-causal features mistakenly influence predictions; (2) Attribute Imbalance certain attribute values appear more frequently than oth- ers; (3) Class Imbalance some labels are significantly underrepresented: and (4) Attribute Generalization mod- els encounter previously unseen attribute values test time. These distinct shift types highlight the diverse challenges of subpopulation robustness and motivate approaches aimed at improving worst-group performance.
Subpopulation shift, characterized by a disparity in subpopulation distribution between the training and target datasets, can significantly degrade the performance of machine learning models. Current solutions to subpopulation shift involve modify- ing empirical risk minimization with re- weighting strategies to improve generalization. This strategy relies on assumptions about the number and na- ture of subpopulations and annotations on group membership, which are unavailable for many real- world datasets. Instead, we propose using an en- semble of diverse classifiers to adaptively capture risk associated with subpopulations. Given fea- ture extractor network, replace its standard linear classification layer with a mixture of proto- typical 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 subpopu- lation shift, our method of Diverse Prototypical Ensembles (DPFss) often outperforms the prior state-of-the-aad in worst-group accuracy. The code is available at https: / httt:::://gtttttt /tggttttbbb.co// com/ o2802222802220 dpw.sspppppppupbpppppp
Naive training using empirical risk minimization (ERM) can result in classifiers achieving good training loss but not generalizing, performing poorly in challenging or un- derrepresented subpopulations (Sagawa et al., 2019; San- turkar et al., 2021). Such failures can be catastrophic in performance-critical real-world applications such as medical diagnostics (Oakden-Ran) Rayner et al., 2020), autonomous driv- ing (Yu et al., 2020), and insurance risk assessment (Bood- hun & Jayahalan, 2018). A commonly cited example is net- works that learn to recognize pneumonia relying on hospital- specific meta-tokens on X-ray scans due to their common co-occurrence in training data (Zech et al., 2018), thus strug- gling to generalize to new data without the tags.
Methods to mitigate subpopulation shift focus on sampling- based (Idrissi et al., 2022; Kirichenko et al., 2022) or loss- based (Michel al., 2022; Han et al., 2022: Sagawa et al., 2020) re-weighting strategies, such as up-weighting minor- ity subpopulations to encourage the model to learn deci- sion boundaries that adequately classify each subpopulation group. But these often require subgroup annotations (Rud- ner et al., 2024; Kirichenko et al., 2022), which are rarely available or sufficiently granular in real-world datasets, or explicit identification of minority groups (Liu al., 2021; Zhang et al., 2022), which significantly increases complex- ity and training time while struggling generalize to unseen subgroups (Yang et al., 2023; Zhang et al., 2022).
1. Introduction
The performance of machine learning models is known to degrade substantially in the presence of distribution shifts be- tween training and deployment (Koh et al., 2021). One com- mon form of distribution shift is subpopulation shift ( (Yang al., 2023), where the proportions of subgroups vary be-
'Deparmment of Electrical and Computer Engineoring, Uni- versity of British Columbia, Vancouver, Canada VVector Insti- tute, Toronto, Canada "School of Computing, Queen's 'een's University, Kingston, Canada "Department of Computer Science, University of Toronto, Toronto, Canada. Correspondence to: Minh Nguyen Nhat To <mtrel student.ubc.ca>>.
We propose Diversified Prototypical Ensembles (DPE). Sub- population shifts results in model degradation because a single classifier typically focuses on the majority classes or subgroups over the minority ones. By turning to ensembles we can capture multiple different decision boundaries. But
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift
Cllsrs
Class K
Task: Binary Classification "Landbisd" "Waterbled'
Class Class 1
Cleas
4
Clues I
Class 1: "Landbard"
"Landbird" (Y=
Prsstype
Proteepe I "Ra] Velle--oeennneee Accested TerdBon Nee Water"
Feature Extractor
Protetgge
Ansaleot (Uuiohelldd) Sulgroup Lonahilds on Land
Postatype
Protetope 3
Amplesi (Ualideelled) Sulgerugs Lasallad Waier
"W "Wateebird" (Y'
Protetype N
Anpaltse (Cndabbtthed) Silggroge Waterbiedd in Water
Protetype XV
Legend
Tarse descriptioss Lave bocn uauurally canotuled to dimorommtte hane DPP Joensi adhmnt xabepoupe within each clsss.
mactimioes
Cleer Ointroii
Parther bum Oweewal
between prototypes in each class to identify relevant subgrosge within the cloce
Implleat (Undnbelhai)) Solgruuge Waterliiid Latel
Figure 1. High-level overview of our method. (1) Binary classification with implicit (unannotated) subgroups. We aim to natively detect and correct subpopulation shifts without prior subgroup knowledge. (2) Given a frozen feature extractor, f(.), we train (3) an ensemble of N prototype classifiers for each of the K classes to identify distinct sub-groups. These classifiers are trained using Lis (Equation 4) to maximize prototype diversity, ensuring robust subpopulation capture within each class. (4) A low-dimensional projection of the centroids and proximal images for class "Landbird" in Waterbirds. The learned centroids for each ensemble member reveal unique latent subpopulations. Points closest to each centroid appear blue, while points farther away are in red. The closest few points shown dark blue, with corresponding images visualized in (5). (5) Visualization confirms DPE's ability capture salient subgroups. We have manually annotated the theme associated with each learned prototype centroid. The closest points to each centroid exhibit thematic consistency, aligning with implicit data subgroups (e.g., birds "on land" vs. "in water").
a naive algorithm for learning ensembles may not necessar- ily encourage each member capture different decision boundary. We take inspiration from recent work studying out of distribution detection and improving generalization (Ginsberg et al., 2022; Pagliardini et al., 2022) by encour- aging diversity among members of an ensembles. If ERM encourages focusing on the majority class' decision bound- ary, explicitly encouraging diversity could encourage sub- sequent members of the ensemble to capture the different decision boundaries corresponding to subgroups even when labels are unavailable or the number, identity distribution of the subpopulations is unknown. Given a feature extractor obtained from standard ERM training, we replace its classifi- cation head with an ensemble of prototype classifiers. Each new prototype serves to predict the class of nearby points in the latent space. An inter-proon piter-prototype similarity loss is used to promote diversity among the ensemble members, resulting in the discovery of many different decision boundaries, each better suited to different subpopulations. This approach yields a model that is robust to subpopulation shifts, as it can adapt to diverse data distributions and maintain classi- fication accuracy across a broad range of subpopulations. Our primary contributions are summarized as follows:
discover and classify subpopulations in the data. Our method improves robustness to subpopulation shifts, even under unknown subgroup annotations, subpopula- tion number, identity, or distribution.
Our solution replaces the linear layer of trained network with the Diversified Prototypical Ensemble (DPE), a collection of distance-based classifiers that incorporate explicit diversification through a loss term and sampling strategy.
We empirically validate DPE using nine real-world datasets proposed in (Yang et al., 2023) to assess ro- bustness against different types of subpopulation shifts. Our results show DPE's superior performance over prior state-of-the-art methods, including in challenging cases like attribute generalization and imbalance.
2. Related Work
Subpopulation Shift Sugiyama al. (2008) character- ized the effects of the covariate shift on a variety of per- formance metrics. Rabanser et al. (2019); Li et al. (2021) showed how metrics changed relatively among the sub- groups of the data under the covariate shift. Yang et al. (2023) created a unified framework for the analysis of the performance of models under various types of subpopula- tion shift.
We propose a novel differentiable end-to-end solution to subpopulation shift based on the idea diversified ensemble, training a collection of diversified predictors
2

---

Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift
Several methods have been proposed to mitigate the ef- fects of spurious corrllations in DL. models. By optimizing for the worst-performing subgroup, Group Distributionally Robust Optimization (gDRO) (Sagawa et al., 2019; 2020) forces the model to learn features that are robust predictive across all subgroups, while UnLearning from Experience (ULE) (Mitchell et al., 2024) trains a student model with no constraints to pursue the spurious correlations in the data while teacher model is trained to solve the same prob- lem while avoiding the student's mistakes. Hierarchical methods tackle the same problem by defining structured feature space via predefined label taxonomies (Mukherjee et al., 2023) to improve worst-group generalization under structured shift. Liang et al. (2022) propose a prototyper pe- based incremental learning method that sequentially adapts classifiers to new subpopulations using margin-enforce loss, aiming to balance acquisition and forgetting in the presence of subpopulation shift. Just train twice (JTT) (Liu et al., 2021) trains the model twice, with the second stage mini- mizing the loss over training examples from a resampled dataset that are misclassified at the end of the first stage. Idrissi al. (2022); Deng et al. (2024) highlight the effects of simple data balancing and subgroup-balarde balanced sampling on worst-group accuracy, showing that simple reweighting and resampling can achieve state-of-the-ar performance on most benchmarks. Unlike this prior work, which often relies on explicit subgroup annotations, our method automatically discovers subgroups to accommodate subpopulation shift without prior knowledge of subgroup identities.
their class. Prototypical networks (Snell et al., 2017) ap- proximated latent space using a neural network, where classification was performed using distances to latent proto- types of each class. The automated discovery of prototype latent spaces has found success in supervised vision (Yang et al., 2018) and image segmentation (Dong & Xing, 2018), Prototype-based classification has proven effective in few- shot learning (Snell et al., 2017) and robustness against shortcut learning (Wei et al., 2023) by leveraging distance- based representations. Their results show promising applica- tions of prototypical classifiers in addressing subpopulation shift, since most successful methods rely on the limited availability of balanced group- -anottand data. Our method extends these findings by training multiple prototypes per class, allowing the ensemble to adapt to heterogeneous sub- populations within each category.
Ensemble Diversity Diversification of ensemble mem- bers is thought improve the robustness and generalization of ensemble learning methods (Fort et al., 2019). Several studies have explored methods to promote diversity and extract hidden patterns in the latent space, such as varying training data subsets (bagging and boosting) (Chawla et al., 2003; Seiffert et al., 2009; Zhang et al., 2019), or using regularizers (Xie et al., 2017: Xie, 2015). Diversification through disambiguation (DivDis) (Lee et al., 2022) aims to enhance ensemble diversity by resolving uncertainties that arise from ambiguous data representations. DivDis intro- duces heterogeneity by encouraging models to specialize in different interpretations of overlapping or uncertain in- stances. D-BAT (Pagliardini et al., 2023) learns a diverse ensemble of models by pushing the models to agree on the source distribution while disagreeing on dittr--distribution inputs. Gating networks (Riquelme et al., 2021; Zhou et al., 2022) determine which experts in mixture-of-experts frame- works to activate for a given input, thus promoting diversity by encouraging specialization in different regions of the input space. Our approach explicitly enforces this diversity within a prototypical ensemble, ensuring that each classifier captures complementary subpopulation-spee sppeiffc features, leading improved worst-group accuracy.
Feature Learning Under Subpopulation Shift Geirhos et al. (2020) investigate when and how neural networks learn spurious features that potentially cause performance degradation. Xing et al. (2021); Salazar et al. (2021); Tian et al. (2022) either implicitly or explicitly learn features while optimizing for various faimess metrics. Kirichenko et al. (2022) demonstrated that core information can be ex- tracted from feature representations learned by standard ERM even when spurious correlation exists in training data. Izmailov et al. (2022) explores last-layer retraining by introducing deep feature reweighting (DFR), further high- lighting that ERM-learned features are competitive with those from group robustness methods DFR achieves state- of-the-art art results on many vision and NLP benchmarks. Qiu et al. (2023); LaBonte et al. (2024) provide ablation studies of the last-layer retraining paradigm and propose different reweighting schemes to improve group robustness and op- timize execution time. DPE builds on these insights by freezing the feature extractor and replacing the ERM -based classifiers with distance-based classifiers to effectively cap- ture subpopulation structures.
3. Subpopulation Prototypical Ensemble
To minimize training loss, a learned decision boundary often exploits features that correlate with the ma- jority subpopulations but fail generalize. For example, classifier that distinguishes "cow" from "camel" using background cues such as "grassy" versus "sandy" performs poorly on minority cases, such as cows in desert environ- ments. In contrast, forcing the model to consider multiple, diverse decision rules- each relying on distinct feature sub- sets cets-increases the likelihood of discovering a rule that generalizes across subpopulations (e.g., "hump" versus "no
Motivation
Prototypical Networks and Representations Prototype methods use data prototypes as representatives values in
3

---

Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift
Attribute 1
Attribute 3
Attribute
Oc
Class Class 2
Training data
(b) Single prototype classifier
(c) Size-2 prototypical ensemble
(d) Size-3 prototypical ensemble
Figure 2. Motivation of DPE. (a) The synthetic training data consists of classes, with major subgroups containing Attribute and minority subgroups containing Attributes and 3. Training single model on the entire dataset lcads to suboptimal decision boundaries, focusing primarily the major subgroups; (b, c, the number of models in the prototypical ensemble increases, where each member trained to classify based on a distinct attribute, decision boundaries become more refined, improving classification across subpopulations.
hump," and "beige" versus "spotty").
set of data. In our implementation, following similar ap- proaches in the literature, we use a subset of the validation set (Izmailov et al., 2022; Kirichenko et al., 2022), and for robustness, we use a class-balanced subset of these data. Crucially, unlike many existing methods, our approach does not require subgroup attribute annotations at this stage, al- though if such annotations available, we recommend using them to construct stage, which functionally serves to combine our more robust DPE (described in the following sections) with classifier retraining methods like (Kang et al., 2020; Kirichenko et al., 2022).
This intuition serves as the basis for our Diversified Proto- typical Ensemble (DPE). Given a set of features, DPE learns an ensemble distance- based prototypical classifiers, each providing plausible decision rule while exploiting differ- features than other members of the ensemble. We argue that even when one decision rule fails on a given subpopu- lation, other ensemble members are likely to succeed. As result, the ensemble exhibits improved generalization un- der subpopulation shift (Fig. 2). We choose prototypical classifiers as base learners due to their ability to preserve feature space geometry the limited-data regime compared ERM (Snell al., 2017). This approach augments con- temporary methods of mitigating subpopulation shift that require re-training the classifier using small, subgroup- annotated subset of the validation set (Yang et al., 2023).
cutribute-balanee subset for this
Prototype Classifier Given feature extractor f TR. - Rd and collection of K classes, prototype classifier de- fines of learnable prototypes (p(*) 1,... K} Rd. For an input X, the classifier computes the probabil- ity of label y based on the distance between the extracted feature s(X) and each prototype.
Our method proceeds in two stages. First, we train a general- purpose backbone feature extractor using ERM on the train- ing set. Next, we select a subset of data from the validation (Izmailov et al., 2022), and train DPE on this subset using the fixed features extracted by the backbone. Our DPE ensemble is characterized by the use of distance-based prototypical classifiers, explicit (loss-based) diversity regu- larization, and implicit (sampling-bass) based) diversity regular- ization.
exp (-D(f(X),
P()|X)
(I)
ΣK i= exp D(F((X)
where D Rd x Rd + R is scaled Euclidean distance between normalized vectors, as in Macêdo & Ludermir (2021),
( |d,|
y
Two-Stage Training The first step our pipeline entails training a feature extractor f I + Rd map raw inputs into a lower-dimensional representation space. We train our feature extractor convergence the full training set, fol- lowing prior results that classical ERM suffices produce strong features despite subpopulation shift (Izmailov et al., 2022; Kirichenko et al., 2022).
|
|
where d, is a learnable scaling factor. Then, the loss for each data-label pair (X, y) is,
exp D(fo(X)), (())
C(X,)) -log
ΣK exp D(fo(X) ((fo(x))((0)/ p()))))
(2)
We then freeze ff and and train our Diversified Prototypi- cal Ensemble (DPE) classifier using a small, held-out sub-
where 7 is a temperature hyperparameter. We initialize pro- totypes randomly with mean μ = and standard deviation

---

Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift
0.01, and train them by minimizing Equation 2 using stochastic gradient descent.
4. Experiments
4.1. Datasets
We conduct comprchensive experiments across nine real- world datasets spanning multiple domains to evaluate the ro- bustness of DPE against subpopulation shift. These datasets represent diverse challenges, including spurious correla- tions, attribute and class imbalances, and attribute gener- alization, all common sources of subpopulation shift. The datasets are chosen based on their ability to test various robustness aspects in machine learning models, particularly under settings where attribute information may unknown or imbalanced. Specifically. the evaluation includes WA- TERBIRDS (Wah et al., 2011), CELEBA (Liu al., 2015), METASHIFT (Liang & Zou, 2022), IMAGENETBG (Xiao et al., 2021), NICO++ (Zhang et al., 2023), LIVING 7 (San- turkar et al., 2021), CHEXPERT (Irvin et al., 2019), CIVIL- COMMENTS (Borkan et al., 2019), and MULTINLI (Schuh- mann et al., 2022). We use the same training-validaationnnns splits given by (Yang al.. 2023). More details of all datasets are provided in Appendix A.I.
Prototypical Ensemble Rather than a single prototype per class, our method uses an ensemble of N prototypes per class, yielding collection {p) (8) " jim jml,.., This
ensemble classifier implements the following classification
rule to aggregate predictions from the N prototype-based classifiers:
N
arg
max (I.......
1
Σ
(ulx),
(3)
K}
N j=1
where (*() (v(X) is the prediction from the jth ensemble member under Equation 1.
To encourage diverse decision rules across ensemble members, we employ two prototype diversification strategies: an explicit inter-prototype simi- larity (IPS) loss and implicit diversification via bootstrap aggregation. Without these, naïvely training each ensemble member may lead to redundant decision boundaries between members of the ensemble.
Ensemble Diversification
4.2. Attribute Availability
Attribute availability significantly affects model's abil- ity handle subpopulation shift. Following recent works (Kirichenko et al., 2022; Rudner et al., 2024), we consider scenario where the group-annotated validation set is available and can be used for model selection, hyper- parameters tuning, and re-training the classifier head when freezing the feature extractor. Yang et al. (2023) showed that that most methods for robust learning with subpopulation shift greatly benefit from access small set of group- annotated data to improve performance underrepresented groups. However, in many real-world scenarios, attribute annotations are not available during training or validation. In such cases, models must rely on the inherent structure of the data to handle subpopulation shift.
The IPS loss decorrelates the representations of different prototypes of the same class. For the n' 'tth ensemble member, this loss is,
K
Krl (b) P)
)1)) n d
LIPS
((
(4)
where 1 denotes the indicator function, and ( denotes the Euclidean inner product. In Liss, note that terms are scaled by n and d, the number of ensemble members and embedding dimensions, respectively. At each ensemble stage n, we simultaneously optimize all prototypes asso- ciated with that stage {p% 1==1.... while freezing the prototypes that have been optimized in previous stages, (k) |==1.... w We optimize each prototype via stochastic gradient descent on the sum of Lips and ((X.y) (Equation 2).
DPE is designed to perform effectively in both settings. We evaluate DPE's capability to identify and adapt to potential subpopulations based on inherent data distribution alone in the absence of explicit subgroup annotations, and its efficacy in utilizing the subgroup annotations to specifically tailor the prototype representations on well-defined subpopula- tions. Specifically, we use all the aforementioned datasets in experiments without subgroup annotations, and only WATERBIRDS, CELEBA, CIVII. COVII.COMMNTS, MULTINLI, METASFHIT, and CHEXPERT in experiments with subgroup annotations.
In parallel, we apply bootstrap aggregation by training each ensemble member on different class-balanced subset of the validation data. These random subsets expose each prototype to slightly different distributions, implicitly en- couraging diversity in the learned decision boundaries. The ensemble is thus trained end-to-end by sequentially soly- ing the joint prototype and IPS loss for each new ensemble member.
4.3. Baselines
To rigorously evaluate the effectiveness of DPE in handling subpopulation shift, we compare its performance against
5

---

Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift
Table 1. Woost-group accuracy (WGA) the test without subgroup annotations. The top half reproduces baselines from Subpop- Bench (Yang et al., 2023) using the same ERM backbone, while the bottom half includes results from original papers and our method applied to stronger ERM backbone. DPE denotes our diversified prototypical ensemble. Best and second best values are bolded within each half of the table. The best result per group is underlined and bolded, the second-best is in bold. "-" indicates results not reported.
Algorithm
WATERBIRDS
CELEBA
CIVILCOMMETSS
MULTILII
METASHIFT
CHEXPERT
IMAGENETBG
NICO++
LIVING17
ERM
69.1
57.6.0.s
63.2+1.2
66.4
82 12..
41.7
76.8
35.0a4.1
48.0a1.
CRT
76.320
69.620.7
67.820.3
65.4102
83.1
74.620.
78.2
33.3200
ReWeightCRT
76.3
70.7=0.s
64.7=0.2
25.2+0.2
85.1
75.1=0.2
77.5
33.320.0
DFR
89.0=0.2
73.720.8
64.4+0.1
63.8+0.0
81.4+0.1
75.8=0.3
74.4=1.8
38.0
ERM DPE
91.0
81.940.2
69.9
69.3-008
84.1
87.9
50.0
54.0
ERM
77.923.
66.5222
69,4
66, 66.57
80.020.0
75.620.4
86.4
33.3200
53.320.9
RWY
1
82,9
67 6
68. 08.0.11.
AFR
1..4
82.020.5
68.7+0.6
73.4=0.6
ERM DPE
94.1=0.2
84.6
68,9
70.9
83.6
76,8
88.1=0.7
50.0
63.0
baselines and current sfate-of-the-an methods for subpopula- tion shift robustness, including Empirical Risk Minimization (ERM), Classifier Re-train (CRT, ReweightCRT) (Kang et al., 2020), Deep Feature Reweighting (DFR) (Kirichenko et al., 2022), Just Train Twice (JTT) (Liu et al., 2021) and Connect---Contrast (CnC) (Zhang et al., 2022), RWY (Idrissi et al., 2022), Automatic Feature Reweight- ing (AFR) (Qiu et al., 2023), and Group-Aware Priors (GAP) (Rudner et al., 2024). See more details in Ap- pendix A.3.
5. Results and Discussion
5.1. DPE Achieves Strong Subpopulation Shift Robustness Without Subgroup Annotations
When attributes are unknown during both the training and validation phases, the performance of most methods de- grades significantly under subpopulation shift. Table I presents the worst-group accuracy (WGA) across multi- ple datasets, illustrating the challenges that various methods face in adapting to subpopulation shift without subgroup annotations. ERM and ERM" represent naive baselines of performance in that they implement no mechanism of ex- plicitly accommodating subpopulation shift. Specifically, ERM achieves a WGA of only 69.1 1% on the WATERBIRDS dataset and 63.2% on CIVILCOMMENTS. reweighting approaches like CRT and Re ReWeightCRT strug- gle to adapt in these scenarios, as they depend on access to accurate attribute information.
4.4. Implementation
We prepared each dataset in alignment with established benchmarks on subpopulation shift, as outlined in Yang et al. (2023). We adopted pretrained ResNet-50 for image data and BERT for textual data to facilitate a direct com- parison with state-of-the-art benchmarks on subpopulation shift robustness. Hyperparameters related to DPE (e.g., т and ar) are tuned using a held-out subset of the validation set, which split into training and validation folds for tun- ing. For consistency and rigor, both hyperparameter tuning and model selection are based on the worst-group accuracy within the validation fold. After hyperparameter selection, we retrain the prototypical ensemble on the full validation set using the selected hyperparameters.
Additionally,
The proposed method, DPE, significantly outperforms all baseline models handling unknown attributes. As shown in Table DPE consistently achieves higher WGA across various datasets, with average of 73.9%. This is sub- stantial improvement compared to ERM's average of 57.7%, and it also surpasses other robust methods such as DFR and RWY, which achieve 65.2% and 67. 67.5%, respectively. On difficult datasets like CHEXPERT, DPE achieves a WGA of 76.8%, exceeding the performance of CRT (74.6%) and DFR (75.8%). Similarly, on the CELEBA dataset, which includes imbalances in hair color and gender, DPE attains WGA of 84.6%, outperforming other baselines. Overall, the results demonstrate DPE's s capability to implicitly discover subpopulations and capture different aspects of the data.
Each tto-stage training cycle is repeated three times with three different seeds ensure the stability of the method. Compared to the initial training stage implemented in Yang et al. (2023) (reported as ERM), our initial training stage hori-
involves stronger augmentation candom--rop/restte zontal flip) for visual datasets and longer training time for all datasets (similar to Idrissi et al. (2022), reported as ERM"). ). In the second training stage, we trained 15 proto- types sequentially for all datasets. The full implementation is detailed in Appendix B and our code base.
5.2. Subgroup Annotations Further Enhance DPE's Performance in Addressing Subpopulation Shift
When subgroup annotations are available in the validation set, models can leverage this information to improve their robustness, either by using group-balanced subsampling or
6

---

Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift
Spurious Correlations
Attribute Imbalance
Class Imbalance
Attribute Generalization
+36
Worst Worst group improvement
Unknown attributes
+27
Known attributes in validation set
80
+21
+21
50
+15
+13
+17
20
+4
+7
+2
+2
WATERBIRDS
CELEBA
METASHIFT
CiviL COMMENTS
MuuTINI
CHEXPERT
IMAGENerBG
NICO+
LMNN117
Figure 3. Worst-group improvement over ERM when using DPE with and without subgroup annotations.
Table Worst-gn group accuracy (WGA) on the test set with access to attribute annotations. The top half presents SubpopBench-sti baselines and method standard ERM backbone. The bottom half includes recent staeeoffttheaa methods and method applied to a stronger ERM" backbone. DPE refers to our diversified prototypical ensemble. Group Info (Train/Vos (iaii//Va)aa indicates whether group labels required: no group info required, / group info required for hyperparameter tuning, // validation data required for both training and hyperparameter tuning. Best and second-best values within each half are underlined and bolded and bold, respectively. indicates results not reported.
Algorithm
Group Info (Train/Val)
WATERBIRDS
CELEBA
CIVI.COMMENTS
MULTINLI
METASHIFT
CHEXPERT
ERM
XIX
19.1
57.620.8
63.2+1.2
66.4+2.3
82.1
41.7+3.4
CRT
xIV
76.310.8
70.410.4
68.510.0
65.410.1
83.110.0
74.0
ReWeightCRT
XIV
76.320.2
71.
68.220.4
65.3+0.
85.1
73.920.2
X
89.0
86.320.3
66.520.2
63.820.0
81.5220.
75.4
ERM + DPE
x
91.0+0.4
87.7
71.5
74.8
87.9=0.7
ERM"
XIX
77.9+3.0
66.5+2..
69.4+1.2
66.3+0.7
80.0=0.0
75.6
Group DRO
VIV
91.411.1
88,912.3
70.012.0
77.7
RWG
VIV
87.6#1.6
84.3+1.8
72.0=1.9
69 09.6+1.0
JTT
XIV
86.7
81.1
69.3
72.6
CnC
XIV
88.510.3
88,810.9
68.912.1
SSA
XIV
89.020.6
89.8
69.9=2.0
76.620.7
DFR
X
92.940.2
88.3+1.1
70.1
74.7:40.7
GAP (Last Layer)
XIV
93.2+0.2
90.2+0.3
74.3+00
GAP (All Layer)
XIV
93.8
90.2:1033
77.8
ERM" DPE
XIV
94.1 +0.4
90.3+0.7
70.8
75.3+0.0
91.7=1.3
76.0=0.3
reweighting, particularly against underrepresented or chal- lenging subpopulations, Table 2 presents comparison of WGA for several methods across datasets where subgroup annotations are known. Compared to the scenario where attributes are unknown, all models show improved perfor- mance in this setting. For instance, CRT achieves a WGA of 70.4% on CELEBA and 68.5% on demonstrating clear advantage over its performance when attributes unknown. DFR and JTT demonstrate superior performance, achieving WGA of 86.3% and 81.1% CELEBA, respectively. DPE delivers the most substantial improvements across almost all datasets. As seen in Ta- ble 2, DPE achieves the highest average WGA at 83.0%, consistently outperforming other methods. On METASHIF,, DPE reaches WGA of 91.7%, surpassing Re WeighhtUR (85.1%) ()) and DFR (81.5%). For WATERBIRDS, DPE attains WGA of 94,1%, which is higher than GAP (93.8%) and DFR (89.0%). Furthermore, in CIVILCOMMENTS, a dataset
characterized by significant identity-based imbalances, DPE achieves WGA of 70.8%, outperforming CRT and DFR.
5.3. DPE Mitigates Challenging Subpopulation Shift
One of the most critical challenges for machine learning models in subpopulation shift is the ability to handle at- tribute imbalance (AT) and attribute generalization (AG). As shown in the Yang et al. (2023), most existing methods- including well-known algorithms like GroupDRO and JTT- show limited improvement in scenarios involving Al and AG. Yang et al. (2023) highlights that none the existing methods significantly outperform ERM in AG settings, in- dicating their inability cope with the imbalanced/unseen attributes in the test data. In contrast, DPE consistently achieves the highest WGA in both AI (e.g. CHEXPERT, CIVI.COMMENTS) and AG (e.g. NICO++, LIVING17 (Fig. 3, Table 1). DPE's use of a prototypical ensemble

---

Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift
allows each prototype explore substructures in the la- tent space without requiring knowledge about the number of subpopulations or the detection of instances in vulnera- ble subpopulations. The diverse prototypes ensure that the
WATERBIRDS
CELEBA
group aepuraccuuyyyyy (96)
94.0
90.0
91.0
88.0
model does not rely on the most frequent attributes alone
but instead learns a broader representation of the underlying data, resulting in better performance across all subgroups.
92.0
86.0
91.0
Number of prototypes
Number of prototypes
5.4. ERM vs. ERM': Disentangling Backbone Effects
..... No diversification
Random subset
Random subset Proer-Pro Similarity loss
To isolate the impact of our prototypical ensemble from the underlying feature extractor, we compare four variants: ERM, ERM" (with stronger augmentation and longer train- ing), and both combined with DPF. As shown in Tables 2a and 2b, DPE improves worst-group accuracy (WGA) con- sistently regardless of the backbone. For example, on WA- TERBIRDS, WGA increases from 69. 1% (ERM) and 77.9% (ERM)) ) to 91.0% and 94.1% 1% with DPE, respectively. Sim- ilar trends hold across datasets like CELEBA (57.6% + 81.9% with ERM+DPE; 66.5% - 84.6% with +DPE) and MULTINLI (66.4% 69.3%; 66.5% - 70.9%). These consistent gains both with and without subgroup availability indicate that DPE's robustness stems from prototype diver- sification rather than representation quality alone, and that it serves as an effective, modular improvement atop standard training pipelines.
Figure 4. Effect of different ensemble diversification methods on performance with different numbers of ensemble members.
ing on different parts of the data distribution.
In the third scenario, apply both random subset selec- tion and inter-prototype similarity loss (IPS), which ex- plicitly minimizes the overlap between prototypes in latent space. This yields the best overall performance across all datasets, with substantial gains in WGA as the ensemble size increases. IPS ensures that each prototype independently contributes to the overall performance, reducing redundancy and enhancing the ensemble's ability to cover broader range of data distributions with more prototypes.
6. Ablation study
6.2. Prototypical Ensembles Outperform Linear Ensembles
6.1. Prototype Diversification is Critical for Capturing Diverse Subpopulations
We compare the WGA of the prototypical ensemble with that of linear ensemble, focusing on scenarios where both models have the same number of ensemble members and were trained on random subgroup- ballnced subsets. Fig. highlights the performance differences across multiple datasets, demonstrating the clear advantage of the prototyp- ical ensemble over the linear ensemble when dealing with subpopulation shift. The linear ensemble, which consists of multiple independently trained linear classifiers, shows similar performance to the original DFR, which only trained three classifiers as the ensembles. On the other hand, from both Fig. 4 and Fig. 5, the prototypical ensemble shows steady improvement when increasing the ensemble size and consistently outperforms the linear ensemble across all datasets regardless of the attribute availability. More- over, prototype ensemble markedly improves over linear ensemble on attribute imbalance (MULTINLI) and attribute generalization (IMAGENETBG, LIVING 7), the two most challenging types of subpopulation shift.
In this section, we investigate the impact of prototype di- versity strategies on the performance of the prototypical ensemble, when subgroup annotation available, specifi- cally comparing three scenarios: (1) using a fixed subset of data for each prototype, (2) training with random subset for each prototype, and (3) incorporating both random subset selection and inter- prototype similarity loss to maximize di- versity across the ensemble members. Fig. 4 highlights the WGA across different datasets as we increase the number of prototypes under these scenarios.
In the first scenario, where a fixed subset is used for each prototype, the performance initially improves as we increase the number of prototypes. However, this improvement quickly plateaus, suggesting that the prototypes learn re- dundant features when trained on the same subset. Without introducing diversity, adding more prototypes beyond a cer- tain point fails to significantly boost WGA.
The second scenario, involving random subset selection, introduces some level of diversity by ensuring each proto- type is trained on different random subset of data. This strategy leads noticeable improvements (Fig. 4) as the number of prototypes increases, with each prototype focus-
6.3. Effect of Ensemble Size
We analyze the impact of the number of prototypes per class, denoted by N, on worst-group accuracy (WGA). To quantify the performance gain, we define the percentage
8

---

Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift
Spurious Correlations
Attribute Imbalance
Class Imbalance
Attribute Generalization
Worst group accuracy (86
0.8
0.7
WATERBIRDS
CELEBA
METASHIFT
CIviL COMMENTS
MuLNNLI
CHEXPERT
IMAGENETBG
LivNG177
Prototypical ensemble with known/unknown attributes
Linear ensemble with known/unknown attributes
Figure Linear ensemble versus prototypical ensemble with known and unknown attributes. NICO++ is not included the plot since our experiments, the worst-group accuracy of the lincar ensemble this datasct is zero.
WGA, 100, where WGAN denotes the worst-group accuracy when using N prototypes per class. We conduct this ablation on four representa- tive datasets: WATERBIRDS, CELEBA, METASHIFT, and CHEXPERT, evaluated under both known and unknown at- tribute settings.
improvement Av WGA MMMMy-WGM WGA1
6.6. Standard Accuracy and Robustness Trade-off
To evaluate whether robustness gains trade off with overall performance, report average accuracy in Appendix C.6. DPE maintains competitive improved accuracy compared to ERM and ERM", in both known and unknown attribute settings. For example, on METASHIFT, ERM" + DPE im- proves WGA from 80.0% to 83.6% and average accuracy from 93.2% to 93.8%. These results support that DPE en- hances fairness without sacrificing accuracy.
Our results show that increasing N leads to diminishing re- turns beyond a certain point. Specifically, we observe the fol- lowing average improvements: A5 2.4%, A10 3.3%, A15 = 3.7%, A25 = 3.7%, and A4o = 3.7%. These results indicate that WGA saturates around N = 15. and further increases in ensemble size result in minimal additional gain due to overlapping prototypes in latent space. These find- ings empirically support that our use of N 15 balances robustness and computational efficiency.
7. Limitations
A limitation of our approach is some additional complex- ity introduced by using an ensemble and some additional hyperparameter'' introduced by the method; however, this limitation is outweighed by improved performance and greater flexibility, as it does not require identification of sub-populations like JTT (Liu al., 2021) and CnC (Zhang et al., 2022). The method remains computationally efficient as the ensemble is trained only on pre-extracted features, adding approximately two minutes per prototype on most datasets. Another limitation the lack formal theoreti- cal explanation for why prototype diversification improves gorst-group accuracy. To partially address this, we con- ducted an exploratory analysis WATERBIRDE dataset in which a third-party annotator (ChatGPT) examined samples closest to each prototype and found consistent semantic pat- terns such as habitat type and pose, suggesting that DPE may implicitly recover meaningful subgroups despite not using group labels (Appendix C.7). While promising, these findings are qualitative, and developing a theoretical under- standing remains open direction. Finally, similar prior works, our method relies on ERM for feature extraction, which may underperform in low-data or weak-labe settings; future work could explore integrating self-supervised learn- ing to improve generalization in such regimes.
6.4. Sensitivity Hyperparameters
We assess the sensitivity of DPE two core hyperparam- eters: the temperature and the inter-prototype similar- ity loss weight a. Across
{10, 20, 30, 40} and E (10*, 5 10 100 5x 5x000) on WATERBIRDS, METASHIFT, and LIVING 17, performance varies within 1 1-2%, confirm- ing DPE's robustness to hyperparameter tuning. The only exception LIVING17, which shows greater variability due to its fine-grained subpopulation structure. See Ap- pendix C.5 for the full analysis.
6.5. Runtime and Memory Overhead
To quantify computational efficiency, we benchmarked run- time and GPU memory usage on an RTX6000 with ResNet-50 backbone and batch size 1. The results show that increasing the number of prototypes from 15 to 100 leads to modest increase in per-batch latency (from 0,0031s to 0.0045s) and memory (from 0.20 GB to 0.85 GB). These results confirm that DPE is scalable and incurs minimal overhead even with large ensemble sizes.

---

Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift
Acknowledgements
Deng, Y., Yang, Y., Mirzasoleiman, B., and Gu, Q. Robust learning with progressive data expansion against spurious correlation. Advances in Neural Information Processing Systems, 36, 2024.
This work was supported by the Canada CIFAR AI Chair program, the Vector Institute, and the Canada Research Chair Tier I in Medical Informatics. We gratefully acknowl- edge funding from the Canadian Institutes of Health Re- search (CIHR) and the Natural Sciences and Engineering Re- search Council of Canada (NSERC). Additional resources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and the companies sponsoring the Vector Institute.
Dong, N. and Xing, E. P. Few-shot semantic segmentation with prototype learning, In BMVC, volume 3, pp. 4, 2018.
Fort, S., Hu, H., and Lakshminarayanan, B. Deep en- sembles: A loss landscape perspective. arXiv preprint arXiv: 19X1::///2.02777, 2019.
RGK is supported by Canada CIFAR AI Chair and Canada Research Chair Tier II in Computational Medicine.
Geirhos, R., Jacohsen, J.-H., Michaelis, C., Zemel, R., Bren- del, W., Bethge, M., and Wichmann, F. A. Shortcut learn ing in deep neural networks. Nature Machine Intelligence, 2(1 (11):665-673) 2020.
Impact Statement
This work introduces Diverse Prototypical Ensembles (DPE), method designed enhance model robustness under subpopulation shifts without relying on explicit group annotations. By leveraging class-balanced subsampling and inter-protte pioot-prototype diversification, DPE aims to mitigate per- formance disparities across subgroups, promoting fairess in model predictions. The potential societal impact of this work lies in its applicability to real-world scenarios where subgroup labels are unavailable or costly to obtain. DPE's ability improve worst-group accuracy without significant trade-offs for majority groups suggests its utility in domains such healthcare, finance, and education, where equitable model performance is critical. However, it is important to acknowledge that while DPE addresses certain fairness con- cerns, it does not eliminate all biases inherent in data or model architectures. Further research is necessary to assess DPE's effectiveness across diverse datasets and to explore its integration with other fairness-enhancing techniques.
Ginsberg, T., Liang, Z., and Krishnan, R. G. A learning based hypothesis test for harmful covariate shift. arXiv preprint arXiv:2212.02742, 2022.
Han, Z., Liang, Z., Yang, F., Liu, L., Li, L., Bian, Y., Zhao, P., Wu, B., Zhang, C., and Yao, J. Umix: Im- proving importance weighting for subpopulation shift via uncertainty-awant mixup. Advances Neural Informa- tion Processing Systems, 01:37704-37711 2022.
Idrissi, B. Y., Arjovsky, M., Pezeshki, M., and Lopez-Paz, D. Simple data balancing achieves competitive worst- group-accuracy. In Conference on Causal Learning and Reasoning, pp. 336-351. PMLR, 2022.
Irvin, J., Rajpurkar, P., Ko, M., Yu, Y., Ciurea-Ileus, S., Chute, C., Marklund, H., Haghgoo, B., Ball, R., Shpan- skaya, K., et al. Chexpert: A large chest radiograph dataset with uncertainty labels and expert comparison. In Proceedings of the AAAI Conference on Artificial Intelli- gence, volume 33, pp. 590-597, 29,,599,2111 2019.
References
Boodhun, N. and Jayabalan, M. Risk prediction in life insurance industry using supervised learning algorithms. Complex & Intelligent Systems, 145::145515,, 2018.
Izmailov, P., Kirichenko, P., Gruver, N., and Wilson, A. G. On feature learning in the presence of spurious correla- tions. Advances in Neural Information Processing Sys- lems, 11885516--3322 2022.
Borkan, D., Dixon, L., Sorensen, J., Thain, N., and Vasser- man, L. Nuanced metrics for measuring unintended bias with real data for text classification. In Companion Pro- ceedings of the 2019 World Wide Web Conference, pp. 491-500, 2019.
Kang, B., Xie, S., Rohrbach, M., Yan, Z., Gordo, A., Feng, J., and Kalantidis, Y. Decoupling representation and classifier for long-tailed recognition. In International Conference on Learning Representations, 2020.
Chawla, N. V., Lazarevic, A., Hall, L. O., and Bowyer, K. W. Smoteboost: Improving prediction of the minority class in boosting. In Knowledge Discovery in Databuses: PKDD 2003: 7th European Conference on Principles and Practice of Knowledge Discovery in Databases, Cavtat- Dubrovnik, Croatia, September 22-26, 2003. Proceedings 7, pp. 107-119. Springer, 2003.
Kirichenko, P., Izmailov, P., and Wilson, A. G. Last layer re-training is sufficient for robustness to spurious correla- tions. In ICML 2022: Workshop on Spurious Correlations, Invariance and Stability, 2022.
Koh, P. W., Sagawa, S., Marklund, H., Xie, S. M., Zhang, M., Balsubramani, A., Hu, W., Yasunaga M., Phillips,
10

---

Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift
R. L., Gao, I., et al. Wilds: A benchmark of in-the- wild distribution shifts. In International conference machine learning, pp. 5637-5664. PMLR,, 2021.
Oakden-Rayned L.,, Dunnmon, J., Carneioo, G., and Ré, C. Hidden stratification causes clinically meaningful failures in machine learning for medical imaging. In Proceedings of the ACM conference on health, inference, and learning, pp. 151-159, 2020.
LaBonte, T., Muthukumar, V., and Kumar, A. Towards last-layer retraining for group robustness with fewer an- notations. Advances in Neural Information Processing Systems, 36, 2024.
Pagliardini, M., Jaggi, M., Fleuret, F,, and Karimireddy, S. P. Agree to disagree: Diversity through disagreement for better transferability. In ICLR 2022-International Conference on Learning Representations, 2022.
Lee, Y., Yao, H., and Finn, C. Diversify and disambiguate: Learning from underspecified data. In ICML 2022: Work- shop on Spurious Correlations, Invariance and Stability, 2022.
Pagliardini, M., Jaggi. M., Fleuret, F., and Karimireddy, S. P. Agree to disagree: Diversity through disagreement for better transferability. In The Eleventh International Conference on Learning Representations, 2023.
Li, M., Namkoong, H., and Xia, S. Evaluating model per- formance under worst-case subpopulations. Advances in Neural Information Processing Systems, 34:17325- 17334, 2021.
Qiu, S., Potapczynski, A., Izmailov, P., and Wilson, A. G. Simple and fast group robustness by automatic feature reweighting. In International Conference on Machine Learning, pp. 28448 28467. PMLR, 2023.
Liang, M., Zhou, J., Wei, W., and Wu, Y. Balancing between forgetting and acquisition in incremental subpopulation learning. In European Conference on Computer Vision (ECCV), pp. 364-380. Springer, 2022.
Rabanser, S., Günnemann, S., and Lipton, Z. Failing loudly: An empirical study of methods for detecting dataset shift. Advances in Neural Information Processing Systems, 32, 2019.
Liang, W. and Zou, Metashift: A dataset of datasets for evaluating contextual distribution shifts and training conflicts. arXiv preprint arXiv:2202.06523, 2022.
Riquelme, C., Puigcerver, J., Mustafa, B., Neumann, M., Jenatton, R., Susano Pinto, A., Keysers, D., and Houlsby, N. Scaling vision with sparse mixture of experts. Ad- vances in Newral Information Processing Systems, 34: 8583-8595, 2021.
Liu, E. Z.,, Haghgoo, B., Chen, A. S., Raghunathan, A., Koh, P. W., Sagawa, S., Liang, P., and Finn, C. Just train twice: Improving group robustness without training group information. International Conference on Machine Learning, pp. 6781-6792. PMLR, 2021.
Rudner, T. G., Zhang, Y. Wilson, A. G., and Kempe, J. Mind the gap: Improving robustness subpopulation shifts with group-aware priors. In International Confer- ence on Artificial Intelligence and Statistics, pp. 127-135. PMLR, 2024.
Liu, Z., Luo, P., Wang, X., and Tang, X. Deep learning face attributes in the wild. In Proceedings of the IEEE International Conference on Computer Vision, pp. 3730- 3738, 2015.
Sagawa, S., Koh, P. W., Hashimoto, T. B., and Liang, P. Distributionally robust neural networks. In International Conference on Learning Representations, 2019.
Macêdo, D. and Ludermir, T. Enhanced isotropy maxi- mization loss: Seamless and high-perfore perfoorrfommece out-of- distribution detection simply replacing the softmax loss. arXiv preprint ar 105 2021.
Sagawa, S., Raghunathan, A., Koh, P. W., and Liang, P. An investigation of why overparamettrizzation exacerbates spurious correlations. In International Conference PMLR, 2020.
Michel, P., Hashimoto, T., and Neubig. G. Distribution- ally robust models with parametric likelihood ratios. In International Conference on Learning Representations, 2022.
Machine Learning, pp.
Salazar, R., Neutatz, F., and Abedjan, Z. Automated feature engineering for algorithmic fairness. Proceedings of the VLDB Endowment. 14(9):1694-1700 2021.
Mitchell, J., del Rincón, J. M., and McLaughlin, N. Un- learning from experience avoid spurious correlations. arXiv preprint arXiv:2409.02792, 2024.
Santurkar, S., Tsipras, D., and Madry, Breeds: Bench- marks for subpopulation shift. In International Confer- ence on Learning Representations, 2021.
Mukherjee, A., Garg, I., and Roy, K. Encoding hierarchical information neural networks helps in subpopulation shift. IEEE Transactions on Artificial Intelligence, 5(2): 827-838, 2023.
Schuhmann, C., Beaumont, R., Vencu, R., Gordon, C.,
Wightman, R., Cherti, M., Coombes, T., Katta, A., Mullis, C., Wortsman, M., et al. 5b: open large-scale

---

Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift
dataset for training next generation image-text models. Advances in Neural Information Processing Systems, 35: 25278-25294, 2022.
Yang, Y., Zhang, H., Katabi, D., and Ghassemi, M. Change is hard: closer look at subpopulation shift. In Proceed- ings of the 40th International Conference Machine 2023.
Learning, pp.
Seiffert, C., Khoshgoftaar, T. M., Van Hulse, J., and Napoli- tano, A. Rusboost: hybrid approach to alleviating class imbalance. IEEE Transactions on Systems, man, and cybernetics-part A: systems and humans, 40(1):185-197, 2009.
Yu, F., Chen, H., Wang. X., Xian, W., Chen, Y., Liu, F., Mad- havan, V., and Darrell, T. Bdd100k: A diverse driving dataset for heterogeneous multitask learning. In Proceed- ings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 2636-2645, 2020.
Snell, J., Swersky, K., and Zemel, R. Prototypical networks for few-shot learning. Advances in neural information processing systems, 30, 2017.
Zech, R., Badgeley, M. A., Liu, M., Costa, A. B., Ti- tano, J. J., and Oermann, E. K. Variable generalization performance of a deep learning model to detect pneumo- nia in chest radiographs: a cross-sectional study, PLoS medicine, 1002668,, 2018.
Sugiyama, M., Suzuki, T., Nakajima, S., Kashima, H., Von Bünau, P., and Kawanabe, M. Direct importance estimation for covariate shift adaptation. Annals of the Institute of Statistical Mathematics, 60:699-746, 2008.
Zhang, M., Sohoni, N. S., Zhang, H. R., Finn, C.,, and Re, C. Correct-n-contrast a contrastive approach for improving robustness to spurious correlations. In International Con- ference on Machine Learning, pp. 2022.
Tian, H., Zhu, T., Liu, W., and Zhou, W. Image fairness in deep learning: problems, models, and challenges. Neu- ral Computing and Applications, 34(15):2275--2289,, 2022.
PMLR,
Zhang, W., Ramezani, R., and Naeim, A. Wothoost: Weighted oversampling technique in boosting for imbal- anced learning. In 2019 IEEE International Conference on Big Data (Big data), pp. 2523-2531. IEEE, 2019.
Wah, C., Branson, S., Welinder, P., Perona, P., and Belongie, S. The caltech-ucsd birds-200-2011 dataset. 2011.
Wei, Y., Ye, J., Huang, Z., Zhang, J., and Shan, H. Online prototype learning for online continual learning. In Pro- ceedings of the IEEE/CVF International Conference on Computer Vision, pp. 18764-18774, 2023.
Zhang, X., He, Y., Xu, R., Yu, H., Shen, Z., and Cui, P. Nico++: Towards better benchmarking for domain gener- alization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 16036- 16047, 2023.
Xiao, K., Engstrom, L., Ilyas, A., and Madry, A. Noise or signal: The role of image backgrounds in object recogni- tion. In International Conference on Learning Represen- tations, 2021.
Zhou, Y., Lei, T., Liu, H., Du, N., Huang, Y., Zhao, V., Dai, A. M., Le, Q. V., Laudon, J., et al. Mixture-of- experts with expert choice routing. Advances in Neural Information Processing Systems, 35:7103-7114, 2022.
Xie, P. Learning compact and effective distance metrics with diversity regularization. In Machine Learning and Knowledge Discovery in Databases: European Confer- ence, ECML PKDD 2015, Porto, Portugal, September 2015, Proceedings, Part I 15, pp. 610-624. Springer, 2015.
Xie, P.. Singh, A., and Xing, E. P. Uncorrelation and even- regularizer. In Interna-
ness: a new tional Conference on Machine Learning, pp. 3811-3820. PMLR, 2017.
Xing, X., Liu, H., Chen, C., and Li, J. Fairness-aware unsupervised feature selection. In Proceedings of the 30th ACM International Conference on Information & Knowledge Management, pp. 1548-35522 2021.
Yang, H.-M., Zhang, X.-Y., Yin, F., and Liu, C.-L. Robust classification with convolutional prototype learning. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 3474-3482, 2018
12

---

Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift
A. Subpopulation Shift Benchmark
Table 2. Summary of the CEL.EBA dataset.
A.1. Datasets
Dataset CELEBA
# Attr.
# Cls.
# Tr.
# Val.
Test
WATERBIRDS (Wah et al., 2011) is a popular dataset for studying spurious correlations, where the background (wa- ter land) is often confounded with the target label (water- birds landbirds). The key challenge in this dataset that waterbirds are predominantly seen in water backgrounds, and landbirds are seen in land backgrounds. This creates spurious correlation, making difficult for models trained on such data to generalize to cases where birds appear in un- usual backgrounds (e.g., landbirds in water environments). Models are evaluated on their ability to classify these atypi- cal instances correctly.
2
162770
19867
19962
A
Dark hair Blond hair
Samples & Samples
71629 (44.0%)
(41
Train
Fenal:
22880 (16.156)
1267 (0.9%)
Val.
833S (43.0%)
B276 (41.75)
2874 (14.5%)
182 (0 10.9%)
Sampioo
9767 (48.996)
7535 (37 7%)
2490 (12.4%)
(0.9%)
Subgroup
Table 1. Summary of the WATERBIRDS dataset.
Dataset
Attr.
Cls.
# TT.
Val.
Test
WATERBIRDS
2
4795
1199
5794
Fanale dark
Male with daris hair
Farale with blnd
Male with Hlond hair
Landhird : Waterbrd
Figure Class/Attribute distribution in CELEBA dataset.
Samplee
2408 (73.0%)
Train
Land Whtte
1057 (22,0%)
104 (3.8%)
56 (1.29)
CHEXPERT (Irvin et al., 2019) is large-scale medical dataset consisting of chest radiographs. The dataset presents significant class imbalance problem, with certain rare medical conditions being underrepresented in the training data. In this setting, models must be robust imbalanced subpopulation distributions to perform well across both com- mon and rare conditions. Evaluating models on this dataset demonstrates their ability to handle real-world medical tasks where some conditions may appear infrequently but are crit- ical to correctly diagnose.
Samples
467 (30.9%)
466 (30 (3099))
133 (1.1%
133 (1.11)
Samples
2255 (38.9%)
2255 (38.9 (38.9%)
642 (11.1%)
642 (11.1%)
Subgroup
Landhird on land
Landbird on land
Waterbird on land
Waterbird water
Table Summary the CHEXPERT dataset.
Dataset CHEXPERT
# Attr.
# Cls.
# TTr
Val.
Test
Figure Class/Attribte distribution in WATERBIRDS dataset.
167093
22280
33419
CELEBA (Liu et al., 2015) a large-scale facial attribute dataset, widely used to explore the impact of attribute im- balance. The target task is classify images based on hair color (blond vs. non-blond) with the confounding attribute being gender. Since blond hair is predominantly females, models tend to overfit to this correlation, leading to poor performance on subgroups such as blond males or non- blond females. This dataset is used to test whether models can generalize across imbalanced attribute distributions and not rely solely on spurious associations.
COMMENSS (Borkan et al., 2019) is a dataset col- lected from an online comment moderation platform, where the task is to predict toxicity in comments. This dataset is used to test robustness identity-based subgroup shifts, as certain demographic groups (such
in
or gender) underrepresented in the data, and there is a strong risk of models developing biased predictions. The ability to cor- rectly classify toxic comments across all identity groups is key challenge in this dataset, and it provides a benchmark for testing faimess in machine learning models.

---

Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift
Doy Cat
Train
Posilive/ Negalive
Cutdoor Indoor
Samples * Sanuphe # Sanmppoo
336776
Train
300
205
2013
A Sammple Sumple > Sample
784 134
507 (22.3%)
789 (36.7%)
196
6%6)
6013
Val.
JACE
Val.
114 (32.7))
71
312
7B (2) 135%
33 19 5961
10420
Test
0772
Test
273
DRL PAL
22
96
0.23)
(3.75)
191 (21.9%))
(7.4%)
Subgroup
Subgroup
Cott
Figure Class/Attribute distribution in CHEXPERT dataset. a) Fach color represents one subgroup; b) Example images from the corresponding subgroups.
dog
Figure Class/Attribute distribution in METASHIFT dataset.
Table 4. Summary of the CIvILCOMMENTS dataset.
Dataset
# Attr.
# Cls.
Val.
Test
IMAGENETBG (Xiao et al., 2021) is modified version of ImageNet that emphasizes background noise as a source of variation in object recognition tasks. The challenge here lies in generalizing to new images where the background differs from what was seen during training. This dataset tests models' ability disentangle foreground objects from background information, ensuring robustness to background shifts.
CIVI.COMMENTS
148304
24278
71854
MULTINLI (Schuhmann al., 2022) natural language inference dataset where models must predict the relation- ship between pairs of sentences (entailment, contradiction, or neutral). The presence of linguistic artifacts, such as negation words (e.g.. "no" or "never"), is often correlated with the contradiction label, creating a spurious association. This dataset is used test whether models can avoid relying on such artifacts and instead generalize across more diverse sentence pairs.
Table Summary of the IMAGENETBG dataset.
Dataset
# Attr.
# Cls.
# Tr.
# Val.
# Test
IMAGENETBG
N/A
9
183006
7200
4050
NICO+ (Zhang et al., 2023) is another domain generaliza- tion dataset that contains various object categories with focus on unseen attributes during testing. The key challenge in this dataset is that models must generalize to unseen com- binations of object categories and attributes (e.g., "cat in autumn" or "dog on water"). It is an important benchmark for evaluating models' ability handle attribute generaliza- tion and unseen variations in test data.
Table 5. Summary of the MULTINLI dataset.
Dataset
Attr.
Cls. Tr.
# Val.
Test
MULTINLI
2
206175
82462
123712
METASHIFT (Liang & Zou, 2022) is dataset that intro- duces the concept of contextual distribution shifts. It con- tains various object classes (e.g., cats and dogs) in different contexts (e.g., indoor and outdoor environments). The task to classify objects while being robust to changes in the context in which these objects are seen. This dataset is par- ticularly useful for testing whether models can generalize to unseen combinations of objects and contexts during testing.
Table 8. Summary of the NICO+ dataset.
Dataset
# Attr.
# Cls.
# Tr.
# Val.
# Test
NICO+ 6
60
62657
8726
17483
LIVING (Santurkar et al., 2021) is part of the BREEDS benchmark, a dataset constructed to test hierarchical clas- sification and domain generalization. In this dataset, the task is to classify living organisms into subcategories, and the challenge comes from unseen subclasses the same
Table 6. Summary of the METASHIFT dataset.
Dataset
# Attr.
# Cls.
Tr.
# Val.
# Test
METASHIFT
2
2276
349
874
2

---

Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift
Train
Sumpler
14700 (23.%%
GTUB (ED 750
Val.
Sampley
Dog
Bird
Wheeled vehicle
Test
Samples
2606 115-290
Subgroup
Reptile
Carnivore
Insect
Autura
Dim
Grass
Outdoor
Rock
Water
Figure 6. Attribute distribution in NICO+ dataset. a) Each color represents one subgroup: b) Example images from the correspond- ing subgroups.
Musical instrument
Primate
Fish
BALANCED ACCURACY (BA)
Balanced accuracy is designed to mitigate the effects of class imbalance by averaging the accuracy across all classes. Let C' denote the set of all classes, and De the subset of data belonging to class c E C. The per-class accuracy is defined as:
Figure Example images from each class IMAGENETBG dataset. Attributes are not available for this dataset.
hierarchical level during testing. This dataset assesses mod- els' robustness in generalizing across bierarchical structures, where unseen subcategories must be correctly identified.
Acc.
Σ I(f(z) 9).
[Del
(x,8)ED.
The balanced accuracy is computed as:
Table 9. Summary of the LIVING 17 dataset.
BA
Acc.-
Dataset
Attr.
# Cls.
# Tr.
# Val.
# Test
CEC
LIVING17
N/A
17
39780
4420
1700
Balanced accuracy ensures that each class contributes equally the overall evaluation, regardless of its frequency in the dataset.
A.2. Evaluation metries
WORST-GROUP ACCURACY (WGA)
A.3. Baselines
Worst-group accuracy measures the performance of a model on the subgroup of data where performs the worst. For- mally, let & denote the set of all groups, and let D, represent the subset of data belonging group g g. Define the accuracy on group y as:
Empirical Risk Minimization (ERM) serves the stan- dard training approach without any modifications for robust- ness. ERM optimizes for overall accuracy, often leading to suboptimal performance on underrepresented subpopu- lations, it tends to prioritize the majority groups during training. While ERM is effective balanced datasets, has well-documented limitations in cases of spurious corre- lations or significant subpopulation shifts, making it key point of comparison.
Σ T(((x) =)) (E,))ED,
Acca = [Dgl
where f(x) the model's prediction and II(-) the indicator function. The worst- worst-group group accuracy is then given by:
Classifier Re-train (CRT, ReWeightCRT) (Kang al., 2020) focuses first learning generalizable representa- tions with standard sampling techniques. Then, in the sec- ond stage, the classifier is re- trainided separately with class- balanced sampling or using non non-parane con-aaamettric methods like Nearest Class Mean (NCM) r-normalization adjust
WGA min Accy- DEG
This metric evaluates the robustness of the model to under-
represented challenging groups in the dataset.
3

---

Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift
decision boundaries for better performance across underrep- resented classes
Automatic Feature Reweighting (AFR) (Qiu al., 2023) enhances group robustness without requiring access group labels during training. AFR operates in two stages: first, standard ERM model is trained on the full dataset. Then, AFR retrains only the last layer of the model on a reweight- ing set, automatically giving higher weight to examples where the ERM model underperforms, which typically be- long minority groups.
Deep Feature Reweighting (DFR) (Kirichenko et al., 2022) aims to mitigate the impact of spurious correlations by de- coupling feature learning from classifier training. DFR specifically retrains the classification layer on a balanced subset of the validation data, ensuring that learned features are leveraged in a more robust manner. These approaches offer improved performance subpopulations but rely heavily on having access to validation data with known subgroups.
Group-Aware Priors (GAP) (Rudner et al., 2024) pro- poses a novel method to improve robustness to subpopu- lation shifts by introducing data-driven priors that favor models capable of generalizing well across different sub- groups. GAP requires only a small set of group-labeled data and uses this to construct probabílistic prior distribution over model parameters. The method operates in two steps: first, neural network is trained using ERM; second, the group-aware prior is used to finetune the model, either on the entire network or by retraining just the last layer. By incorporating this group-aware prior, the model places a higher probability on parameters that achieve strong per- formance on worst-case subgroups, significantly improving group robustness.
Just Train Twice (JTT) (Liu et al., 2021) proposes a two- stage approach to improve worst-group accuracy without requiring group annotations during training. In the first stage, JTT trains standard ERM model and identifies the training examples it misclassifies, which often belong to groups affected by spurious correlations. In the second stage, JTT retrains the model by upweighting these misclassified examples to enhance the model's performance on the worst- performing groups. Despite not using group annotations during training, JTT significantly improves worst-group accuracy, achieving performance close to methods that do rely on group annotations, while maintaining competitive average accuracy across various datasets.
Cortect-n-Contrast (CnC) (Zhang et al., 2022) designs a two-stage contrastive learning method to improve robust- ness against spurious correlations without requiring group labels during training. In the first stage, an ERM model is trained infer spurious attribute labels by predicting groups that may correspond to spurious correlations. In the second stage, CnC uses contrastive learning to align the rep- resentations of same-class samples while ignoring spurious attributes. By sampling same-class instances with different spurious attribute predictions as positives, and different- class instances with the same spurious attribute predictions as negatives, CnC ensures that the learned representations focus meaningful, class-specific features rather than spu- rious correlations.
RWY (Idrissi et al., 2022) presents simple yet effec- tive approach to improving worst-group accuracy through basic data balancing techniques such as subsampling and reweighting without requiring subgroup annotations. The method focuses on addressing group and class imbalances in datasets by either subsampling large groups/classes or reweighting samples to ensure balanced mini-batches during training. The results show that these simple data balancing techniques achieve competitive worst-group accuracy com- pared to more complex state-of-the methods. Addition- ally, the study emphasizes that while attribute information crucial for model selection during validation, it is less important during training.

---

Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift
B. Implementation details
The training procedure for DPE framework consists of two stages and implemented with dataset-specific hyperparameters detailed in Tables 10 and 11. Hyperparameters were selected to optimize the WGA on the validation set, ensuring fair evaluation across different datasets.
B.1. Representation Learning
In the first stage, feature extractor trained using ERM with setaset-specific configurations (Table 10). Image datasets such WATERBIRDS, METASHIFT, and IMAGENETBG use the SGD optimizer with learning rate of le-2 and batch size of 128, while text-based datasets such CIVILCOMMENTS and MCLTINLI use BertAdam with learning rate of 1e-4 and batch size of 16. The number of training epochs varies by dataset, ranging from 4 for CIVILCOMMENTS to 300 for WATERBIRDS.
Table 10. Hyperparametees of the representation learning.
Dataset
# Epochs
Optimizer
LR
Batch size
WATERBIRDS
300
SGD
le-2
128
CELEBA
50
SGD
le-2
128
CIVIICOMMENTS
4
BertAdam
1e-4
16
MULTINLI
5
BertAdam
le-4
16
METASHIFT
100
SGD
le-2
128
IMAGENETBG
20
SGD
le-2
128
NICO+
100
SGD
1e-2
128
CHEXPERT
30
SGD
le-2
128
LIVING17
50
SGD
le-2
128
B.2. Prototypical Ensemble Learning
In the second stage, maass/rroup-balancee subsets of the validation set used to train the prototype ensemble. Each ensemble member is trained sequentially with inter-prototype similarity loss (Lips) to encourage diversity among prototypes. The coefficient X for Luss varies by dataset and is listed in Table 11, with values such as 5 x 10 for WATERBIRDS and 1 x 10a for CIVILCOMMENTS. Most datasets, including WATERBIRDS, CELEBA, and CHEXPERT, utilize the SGD optimizer with a batch size of 256, while METASHIFT uses a smaller batch size of 64 due to its smaller dataset size.
Table II. Hyperparameters prototypical ensemble learning.
Dataset
LR
Optimizer
Batch size
A
WATERBIRDS
le-3
SGD
256
5c5
CELEBA
5e-4
SGD
256
5e5
COMILCOMMENTS
le-4
SGD
256
le5
MULTINLI
le-4
SGD
256
le5
METASHIFT
le-2
SGD
64
le5
IMAGENETBG
le-3
SGD
256
5e5
NICO+
le-2
SGD
256
le5
CHEXPERT
le-3
SGD
256
5e5
LIVING17
5e-5
SGD
256
le5

---

Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift
C. Additional Results
C.1. Effect of Diversification Strategies on Ensemble Performance
Figure 7 extends the results shown in Figure 5 of the main paper to four datasets, illustrating the effect of different ensemble diversification strategies on both worst-group accuracy and balanced accuracy as the number of ensemble members increases. Three methods are compared: fixed subset training (no diversification). (2) random subset training, and (3) random subset selection combined with inter-prototype similarity loss. The results show that the combined approach consistently achieves the highest worst-group accuracy across datasets. This underscores the importance of using both explicit and implicit diversification mechanisms to ensure that the ensemble captures broad range of data distributions and subpopulation dynamics.
WATERBIRDS
CELEBA
METASHIFT
CHEXPERT
92.0
Worst group accuracy
76.0
93.0
88.0
75.2
88.0
92.0
86.0
74.4
86.0
91.0
Balanced accuracy (%)
95.6
92.8
94.0
79.5
95.2
92.4
93.5
79.2
94.8
92.0
93.0
79.0
94.4
92.5
10
10
10
10
Number of prototypes
Number of prototypes
Number of prototypes
Number of prototypes
No diversification
Random subset
Random subset Inter-Prototype Similarity loss
Figure 7. Effect of different ensemble diversification methods on performance with different numbers of ensemble members.
Landbird Prototypes
Waterbird Prototypes
Without Diversification
ET
XY:
CR
832
C.2. Prototype Diversity Visualization
Figure presents pairwise cosine similarity matrices for the first five prototypes within the ensemble, visu- alized separately for the "Landbird" and "Waterbird" classes. These matrices quantify the similarity be- tween prototype embeddings, where values closer to O indicate greater diversity. The visualization demon- strates that our inter pooot-prtotype similarity loss effec- tively encourages representation diversity by reduc- ing prototype overlap, enabling different prototypes capture distinct subpopulation characteristics.
0.5
0.0
with Diversification
-0.5
512
635
6s0
621
-1.0
025
428
2000
Figure 8. Pairwise cosine similarity matrices for the first prototypes in the ensemble of the WATERBIRDS dataset.
6

---

Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift
C.3. Runtime and Memory Efficiency
To assess the computational overhead introduced by our prototypical ensemble, we benchmarked DPE on an RTX6000 GPU using ResNet-50 backbone and k=1000 classes, with batch size of 1. Table 12 shows the runtime and memory usage function of the number of prototypes per class. We observe that inference time increases only slightly the number of prototypes increases, from 0.0031s with 15 prototypes to 0.0045s with 100 prototypes. Memory usage grows more noticeably, increasing from 0.2032 GB to 0.8517 GB. In comparison, a standard linear classifier requires 0.0032s per batch and 0.104 GB of memory. These results suggest that DPE introduces minimal runtime overhead while remaining memory-efficient, even when using up to 100 prototypes, making it practical for large-scale applications.
Table 12. Time and memory benchraaring on RTX6000 with k= classes (ResNet-50, batch size 1).
Model Head
# Prototypes
Time per Batch (s)
GPU Memory (GB)
DPE
15
0.0031
0.2032
DPE
20
0.0033
0.2413
DPE
30
0.0033
0.3176
DPE
100
0.0045
0,8517
Linear (Baseline)
N/A
0.0032
0.1040
C.4. Different Types of Ensemble-based Techniques
this ablation study, we compare the performance of three ensemble learning grrategess-Voting. Bagging, and Stacking-to the proposed ensemble approach, evaluating their efficacy handling classification tasks under diverse conditions. The Voting ensemble aggregates predictions from heterogeneous base models, including Logistic Regression, Decision Tree, and Support Vector Machine (SVM), by employing both hard voting (majority rule) and soft voting (weighted probabilities). The Bagging ensemble utilizes Decision Trees as base learners and trains multiple instances on bootstrapped subsets of the training data, reducing variance and improving generalization. In contrast, the Stacking ensemble leverages a meta-learner (Gaussian Naïve Bayes) to integrate predictions from diverse classifiers (Logistic Regression, Decision Tree, and SVM), optimizing the final decision boundary through hierarchical learning. All models were trained on the same input features from the validation set used for training DPE. The group- balanced sampling strategy used for WATERBIRDS and CELEBA, while class-balanced sampling strategy is used for LIVING 17. Predictions were extracted from both the ensemble models and their individual base learners to analyze decision consistency and diversity. The results in this benchmark show that DPE consistently outperforms existing ensemble-bate based approaches under subpopulation shifts.
Table 13. Performance comparison of ensemble methods.
Method
CELERA
WATERBIRDS
LIVING17
BAce
WGA
BAcc
WGA
BAcc
WGA
Voting
92.65
89.61
95.05
92.06
82.88
47.00
Bagging
90.87
81.11
94.01
92.06
75.71
34.00
Stacking
92.47
85.56
95.36
93.66
74.24
36.00
DPE
92.87
90.31
95.61
94.13
87.03
63.00
C.5. Sensitivity Analysis
We evaluate the robustness of our method to two critical hyperparameters: the inverse temperature parameter (1/7) used in the entropic similarity function, and the IPS loss coefficient (a) that controls the strength of prototype diversification. Table 14 presents the worst-group accuracy across multiple values of each hyperparameter, while Table 15 reports the corresponding overall accuracy. We observe that performance is generally stable across a wide range of settings for both hyperparameters. Notably, the method maintains high accuracy on WATERBIRDS and METASHIFT, with only moderate sensitivity observed on the more challenging LIVING17 dataset. These results suggest that the method does not require fine-tuning to achieve strong robustness and generalization across subpopulations.

---

Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift
Table Sensitivity analysis of worst-group accuracy (WGA) the inverse temperature (1 (1/T)
the loss coefficient (a).
Dataset
1/T
Summary
a
Summary
10
20
30
40
Mean
STD
le4
5e4
le5
5e5
Mean
STD
Waterbirds
93.6
93.9
94.1
94.5
94.03
0.38
93.6
94.4
94.1
94.1
94.05
0.33
MetaShift
89.7
91.7
90.8
91.3
90.88
0.87
90.5
91.7
90.7
90.8
90.93
0.53
Living17
63.0
58.7
57.3
55.3
58.58
3.26
63.0
61.7
61.7
62.0
62.10
0.62
Table 15. Sensitivity analysis of overall aceuracy the inverse temperature (1/7) and the IPS loss coefficient (o).
Dataset
1 1/T
Summary
a
Summary
20
40
Mean
STD
le4
5o4
lc5
5c5
Mean
STD
Waterbirds
96.4
96.2
96.0
95.9
96.13
0.22
96.3
95.9
96.0
95.9
96.03
0.19
MetaShift
93.8
93.7
93.9
93.9
93.83
0.10
93.9
93.9
93.7
93.7
93.80
0.12
Living17
87.0
87.2
87.1
86.9
87.05
0.13
87.2
87.0
87.0
87.1
87.08
0.10
C.6. Standard Accuracy
In addition to robustness metrics such as worst-group accuracy, we report the standard (average) accuracy of all evaluated methods across six benchmark datasets in Table 16. These results provide complementary view model performance, reflecting how well each method performs the overall population rather than just the most challenging subgroups. While some methods show strong gains in worst- gros- accuracy, they may trade off slightly in overall accuracy, particularly when subgroup reweighting is involved. Our proposed DPE method achieves competitive or superior average accuracy, especially when combined with stronger ERM" backbone, indicating that robustness improvements do not come at the cost of general performance.
Table 16. Standard (average) accuracy for all methods across datasets. Group Info indicates whether group labels are used for training and/or validation: xix: No group info required, X/V: Group info used for hyperparameter tuning, XIV /: Validation group labels required during training and tuning, vIV: Full group labels required.
Algorithm
Group Info
WATERBIRDS
CELEBA
CIVILCOMMENTS
MULTINLI
METASHIFT
CHEXPERT
ERM
XIX
84.1
95.0+0.1
85,41002
80.9=0.3
91.5+002
88,610.7
CRT
XIV
89.210.1
94.110.11
83.0:10.0
80.2-0.0
91.5.10.0
79.1
ReWeightCRT
XIV
89,410.33
94.2.10.1
83.4.1..0
80.2-0.0
91.3.10.1
79.0.1..0
DFR
XIV
92.210.2
91.2.10.1
81.310.0
80.210.0
90.5110.4
78,910.2
ERM DPE
X
92.5+0.2
89.8+00.2
82.2+0.2
81.3=0.2
91.2++0.1
ERM*
XIX
92.1+0.2
94,0++.22
83.3+1.44
81.9=0.2
93.2+0.11
79.4+0.3
Group DRO
VIV
93.5
92.9
88.9
81.4
RWG
VIV
XIV
93.3
88.0
91.1
78.6
CnC
XIV
90.940.1
89,940.55
81.7+-0.5
SSA
X
92.210.99
92.8.10.1
88.242..0
79.9-1.87
DFR*
X
94.2+00.4
91.3++0.
87.2+0.3
82.1=0.2
GAP (Last Layer)
XIVX
94.6+0.2
91.7++0.2
81.9=0.0
GAP (All Layer)
XIV 95.6+0.11
91.5++0.1
82.5=0.1
ERM* + DPE
XIVX
96.010.1
91.91003
81.6110.2
81.6-0.2
93.8.0.55
79,010.2
C.7. Exploratory Analysis of Prototype-Subgroup Alignment
To gain insight into the semantic structure captured by the Diversified Prototypical Ensemble (DPE), we performed an exploratory analysis on the WATERBIRDS dataset. For each prototype, we retrieved the top-10 closest validation samples and used ChatGPT to interpret the emerging patterns. As shown in Figures and 10, the learned prototypes exhibit consistent alignment with ecologically visually meaningful subpopulations, even though no subgroup labels were used
8

---

Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift
"Diving surface-swimming waterbirds open aquatic environments"
This prototype captures ecologically coherent waterbirds mostly grebes, cormorants, and gulls adapted diving floating The visual consistency plumage contrast, aquatic setting, and makes this prototypical representation of the waterbird class.
"Large-bodied waterbirds in extended flight or lift-off poses" Prototype 2 groups waterbirds broad wingspans, often shown flight preparing takeoff, including albatrosses, gulls, terns. While backgrounds vary (bamboo, docks, ponds), the shared latent subpopulation that reflects kinetic visual mode aquatic birds.
aerodynamic posture silhouette geometry, forming distinct
"Compact-bodioe" diving birds with rounded silhouettes terrestrial ambient-light settings" This prototype groups grebes, auklets, guillemots,
birds morphologically compact, rounded low profiles dark plumage. Despite their aquatic nature, they're often shown mismatched land-based bamboo-heavy suglzesting visually- coherent but ecologlcally-confounde latent cluster.
"Standing gulls upright posture
visual separation"
This prototype forms highly coherent group of California, Glaucous-winged, Heermann's, and Ring-billed Gulls, unified their gray-and-white plumage, yellow bills, and erect standing poses. tolerates wide range of backgrounds, but emphasizes pose uniformity and size. single Hooded Merganser outlier appears, likely due to visual mimicry in shape rather than ecological taxonomic alignment.
Figure 9. Waterbird Prototypes. Each row depicts the top- 10 validation samples closest to one of the prototypes learned for the WATERBIRDS class. Using ChatGPT for cluster interpretation, observe that the prototypes induce structured prototype-subgroup alignment that meaningfully reflects bird morphology. pose, and context: divers and floaters in aquatic settings, large-bodied birds in dynamic flight poses, compact-bodied seabirds shown in terrestrial or bamboo-heavy scenes (an ecologically confounded but visually coherent mode), and upright-standin gulls with consistent visual separation.
during training. In the WATERBIRDS class, prototypes capture distinctions such as aquatic divers, large-bodied birds in dynamic flight, and seabirds in bamboo-heavy or terrestrial contexts. In the LANDBIRDS class, clusters emerge that reflect postural cues, background settings, and potential spurious correlations, such as songbirds appearing in human-associated environments. These findings suggest that DPE implicitly encourages subgroup discovery through diversification, which may contribute to strong sorst-group performance.

---

Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift
"Small yellow warblers and vireos in varied habitats"
This prototype captures highly color-consistent group of small yellow songbirds, largely comprising warblers, vireos, and flycatchers. While visually coherent in terms of plumage and posture, the birds are shown diverse or mismatched backgrounds, suggesting the prototype focuses on foreground semantic structure rather than context. It represents meaningful intra-class mode lancbirds, especially the yellow-dominated cluster.
"Upright-posturee" postured forest and edge-dwelling birds with structured or artificial backdrops This prototype captures vertically oriented birds (especially woodpeckers and similar species) in settings that lack dense greenery -omterryoffen featuring man-made structures, open water, docks, stylized vertical environments (e.g. bamboo). It may encode a shape-surface or pose-context association, blending natural and built environments.
"Forest songbirds in atypical coastal or human-present scenes"
Prototype 4 clusters birds whose ecological identity is forest/shrubland but whose visual context is dominated by beach, water, or people. captures subpopulation prone spurious background correlations, making a key prototype for understanding failure under subpopulation shift.
"Muted-tane songbirds in dense green vegetation"
This prototype captures small passerines - mainly ground-rich vegetation. reflects low-visibility subpopulation that trades high contrast for ecological realism, possibly helping the model generalize across scenes with low texture saliency.
warbley, warblers, vireos, and sparrows that share camouflaged color scheme and are seen in leafy,
Figure 10. Landbird Prototypes. Each shows the top- 10 validation samples closest of the prototypes from Diversified Prototypical Ensemble (DPE) model on the LANDBIRDS class the WATERBIRDS dataset. ChatGPT was used to analyze each prototype's semantic structure. The results reveal that leamed prototypes align with coherent visual ecological subpopulations, for example, small yellow songbirds across diverse backgrounds, upright-postured forest-edge dwellers, forest birds appearing in beach or human-present scenes (highlighting spurious correlations), and muted-tone songbirds in leafy, texture-poor vegetation.
10

---

(0) Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift
D. Algorithm
Our training pipeline (Algorithm 1) consists of two stages: feature extractor training and prototypical ensemble training. In Stage 1, the feature extractor and classification head are trained using ERM on the training data optimize feature representations. In Stage an ensemble of class-specific prototypes is initialized and trained on class-balanced subsets or group-balanced subsets of the validation data, depending the availability of the subgroup annotations. A distance- based loss and an inter-prototype similarity loss are used to update each ensemble member. During inference, class probabilities are computed using the joint decision of the members in the prototypical ensemble.
Algorithm Subpopulation Prototypical Ensemble Diversification
Input :Training data Dumin, Validation data Daat Number of ensemble members N, Diversity weight o, Temperature I,...,
Output : Trained feature extractor f, Ensemble of prototypes (p) (1) N
Stage 1: Train Feature Extractor
Initialize feature extractor f and classification head g for each minibatch (Nuuch, Ybuach) Derain do
Compute logits (( Compute loss Lce E,
E.Bhansh Bhunch log (=2()) Update f and g to minimize LCE
Stage 2: Train Prototypical Ensemble Initialize ensemble of prototypes P = {} for 1 to N do
Sample class-balanced subset ( from Doa Initialize prototypes 1 with p3 ~ NO 0.01²) uhh for each minibatch (X, D do
Compute features f(X) Compute distances D(x, ( using Equation (5) Compute loss C using Equation
j then
Update ensemble P with prototypes (p) ) Compute inter- anter-proootpp similarity loss Lis (Equation (8) using P and ) Set Lioral C Laps
else
Set Lreu C Update prototypes {p) 1 Jim1 minimize Leaa
Update ensemble P with the newly trained prototypes (p) ]
Inference
Given test input X Compute feature x = f(X) for to K
Compute ensemble probability:
exp ( D(x.p, p.,)) 1)
P(y=i[X)= =
k- exp D(2, p2,1)))))
Predict label arg max; Ply x


Reviews:
[{"review_id": "gBu0RwvCYP", "reviewer": "Reviewer_rN4y", "summary": "This paper introduces the Diversified Prototypical Ensemble (DPE) to improve the robustness of machine learning models to subpopulation shifts. It replaces the standard linear classification layer with an ensemble of distance‐based prototypical classifiers. A two-stage training scheme is used: first use ERM to train a feature extractor and then fine-tune the ensemble on a held-out validation set. Empirical evaluations across nine diverse real-world datasets demonstrate that DPE is better in worst-group accuracy than state-of-the-art methods, both when subgroup annotations are available and unavailable.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review. We appreciate your recognition of our well-organized presentation, the strong empirical performance of DPE on worst-group accuracy, and its relevance to prototypical networks and subpopulation shift.  \nYou raised key concerns regarding:\n\n- (1) ablation on prototype count and hyperparameter sensitivity,  \n- (2) missing runtime/resource analysis,  \n- (3) incomplete reporting of average accuracy,  \n- (4) limited evidence linking prototypes to meaningful subpopulations, and  \n- (5) lack of theoretical justification.  \n\nWe address each of these below with new results, analyses, and clarifications. Full tables are provided at https://github.com/anonymous102030411/anon.\n\n### 1. Further Ablation on the Number of Prototypes\n\nWe thank the reviewer for raising this important question. To further investigate the effect of the number of prototypes per class, we extended our ablation study beyond 15 prototypes and now report results up to 40 prototypes per class.  \nLet $\\text{WGA}_N$ denote the worst-group accuracy (WGA) when using $N$ prototypes per class. We compute the percentage improvement over the single-prototype case as:\n\n$\\Delta_N = \\frac{\\text{WGA}_N - \\text{WGA}_1}{\\text{WGA}_1} \\times 100\\%$\n\nWe evaluated this on four representative datasets—Waterbirds, CelebA, Metashift, and CheXpert—under both settings: with and without subgroup annotations. The average percentage improvements are as follows:\n\n$\\Delta_5 = 2.4\\%$, $\\Delta_{10} = 3.3\\%$, $\\Delta_{15} = 3.7\\%$, $\\Delta_{25} = 3.7\\%$, $\\Delta_{40} = 3.7\\%$\n\nThese results show that worst-group accuracy increases rapidly with the number of prototypes up to $N=15$, but plateaus beyond that point, indicating *diminishing returns*. Specifically, increasing from $N=15$ to $N=40$ provides no further gain in WGA.\n\nWe interpret this as empirical evidence that a moderate number of diversified prototypes (e.g., 10–15 per class) is sufficient to capture the key subpopulation structures. Beyond this range, new prototypes tend to overlap in latent space with existing ones, limiting additional benefit. We also note that larger ensembles increase computational and memory costs without proportional gains. These observations justify our choice of using $N=15$ in the main experiments, balancing performance and efficiency. \nRegarding the performance sensitivity to other important hyperparameters in our study, we direct the reviewer to our response to Reviewer [VrPG](https://shorturl.at/t93so) with more extensive results and discussion.\n\n### 2. Computational Complexity Analysis\n\nWe clarify that “complexity” refers to implementation (e.g., added hyperparameters), not compute overhead. DPE adds minimal training cost, as prototypes act like lightweight linear layers. Compared to DFR, we train 15 simple classifiers instead of one. At inference, the overhead is negligible relative to the feature extractor.\n\nTable 1 (See **[Table 1](https://shorturl.at/KekLM)**) compares DPE and DFR in sample throughput and memory usage. Increasing from 15 to 100 prototypes only slightly raises time per batch (0.0031s→0.0045s) and keeps GPU memory below 1 GB—well within typical deep learning budgets. With fewer classes, the relative increase in complexity is even smaller. The complexity analysis will be included in the appendices of the revised version.\n\n### 3. Reporting Average Accuracy\n\nWe appreciate the suggestion to include average accuracy alongside worst-group accuracy. We revised the paper to include average accuracy in Table 6 (see **[Table 6](https://shorturl.at/KQEil)**), demonstrating that DPE remains competitive on average accuracy while prioritizing worst-group robustness.\n\n### 4. Clarifying the Link Between Learned Prototypes and Subpopulations\n\nOur core contribution is demonstrating that prototype diversity reliably improves robustness across a wide range of subpopulation shift benchmarks. We provide the new qualitative visualization showing that learned prototypes naturally cluster semantically related samples and often align with meaningful subgroups, even without explicit supervision (see **[Figure 1 and 2](https://shorturl.at/o6tUI)**). This supports our central claim that diversity in prototype space enables better subgroup coverage. Expanded ablations further support the empirical foundation of our method.\n\n### 5. Theoretical Justification\n\nThe efficacy of DPE is supported primarily through extensive empirical evidence. We provide an intuitive rationale: by diversifying decision boundaries within each class, DPE encourages models to rely on more robust features rather than spurious correlations. We highlight this in the revised discussion and consider a formal theoretical study to be an important direction for future work.\n\n---\n\n**Thank you for taking the time to review our work. If we have answered your questions, then we would appreciate you considering raising your score. If anything is still unclear, we are happy to clarify.**", "claims_and_evidence": "The paper claimed, \"These classifiers are trained using ... to maximize prototype diversity, ensuring robust subpopulation capture within each class\". Although improved worst-group accuracy is provided, the evidence linking each learned prototype to a semantically meaningful subpopulation is somewhat ambiguous.", "methods_and_evaluation_criteria": "For evaluation, the paper mainly focuses on worst group accuracy when compared with other methods. The performance of average accuracy is not completely reported.", "theoretical_claims": "There are no theoretical claims or proofs.", "experimental_designs_or_analyses": "The ablation study on the number of prototypes is somewhat insufficient. Why limit it to less than 15? I am curious what will happen if you further enlarge the number.\n\nThe authors claimed a weakness in their increased complexity, however, they did not conduct experiments to compare the running time, speed, computation resources, etc.", "relation_to_broader_scientific_literature": "The paper relates to work including prototypical networks, ensembling, and subpopulation shifts.", "essential_references_not_discussed": "Not that I am aware of.", "other_strengths_and_weaknesses": "One strength is the paper is well-organized and easy to follow. Some weaknesses include: no theoretical justification, not solid enough evidence to show why worst-group accuracy improved, and some missing ablation studies on important hyperparameters.", "other_comments_or_suggestions": "NA", "questions_for_authors": "NA.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "O4E2EDWAZl", "reviewer": "Reviewer_VrPG", "summary": "This paper tackles the problem of subpopulation shift in machine learning, where the proportions of different subgroups within a dataset change between training and testing. The authors propose a novel method called Diversified Prototypical Ensemble (DPE) to improve robustness to such shifts. DPE combines prototypical classifiers with an ensemble approach and explicit diversification strategies. The core idea is to train multiple prototype classifiers per class, encouraging them to learn different decision boundaries that capture various subpopulations within each class. Diversification is achieved through an inter-prototype similarity loss (IPS) and bootstrap aggregation. The main result is that DPE significantly improves worst-group accuracy (WGA) on real-world datasets, both with and without subgroup annotations during training.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review. We appreciate your recognition of the novelty and relevance of DPE, its effective combination of prototypical networks and ensemble diversification, and its strong empirical performance on worst-group accuracy across standard benchmarks.\n\nYour main concerns include \n- (1) potential confounding due to stronger ERM* training compared to baseline implementations\n- (2) limited hyperparameter sensitivity analysis, particularly for τ, α, and subset size\n- (3) unclear prototype-subgroup alignment.\n\nWe address each of these points below with new results, analyses, and clarifications. Full tables are provided at https://github.com/anonymous102030411/anon.\n\n### 1. Ensuring a Fair Comparison with Prior Art\nWe thank the reviewer for raising this important concern. To isolate the effect of our method (DPE) from stronger training protocols, we retrained all feature extractors using the exact setup from Yang et al. (2023)—matching architecture, augmentation, and training schedule. Applying DPE on top of these retrained backbones (see **[Table 2a](https://shorturl.at/bUCeK)**) confirms that DPE outperforms all reported baselines, showing the gains stem from DPE itself, not the ERM backbone.\n\nWe also report results for ERM* + DPE, which includes stronger augmentation and longer training (see **[Table 2b](https://shorturl.at/3sAov)**). While these features boost performance, **DPE’s advantage remains consistent across both ERM and ERM*** **setups**.\n\nRegarding fairness in comparing ERM* + DPE to prior work, we note that several baselines use enhanced protocols:\n- DFR includes data augmentation.\n- RWG and RWY use extended training and thorough tuning.\n- CRT’s code includes augmentation, though not mentioned in the paper.\n- CnC, SSA, and GAP lack public code, so augmentation use is unclear.\n\nThus, **comparing ERM*** **+ DPE to baseline reports is fair**, as many already benefit from stronger pipelines. In this context, DPE still achieves state-of-the-art worst-group accuracy across nearly all benchmarks.\nTo support transparency, we will include:\n- A clear discussion of ERM vs. ERM* in the paper.\n- A full performance table with both ERM + DPE and ERM* + DPE results.\n\n### 2. Hyperparameter Tuning and Sensitivity\nWe thank the reviewer for pointing out the importance of hyperparameter sensitivity analysis. In our study, hyperparameters related to DPE—specifically the inverse temperature (**1/τ**), and the IPS loss weight (**α**)—are tuned using a held-out subset of the validation set, which is split into training and validation folds for tuning. Once optimal hyperparameters are selected, the prototypical ensemble is trained on the full validation set using these tuned values. Therefore, the subset size is not a hyperparameter in our pipeline.  \nTo address the reviewer’s concern, we conducted a sensitivity analysis on inverse temperature (**1/τ ∈ {10, 20, 30, 40}**) and IPS loss weight (**α ∈ {1e4, 5e4, 1e5, 5e5}**). Results across three datasets (Waterbirds, MetaShift, Living17) are presented in the attached tables (see **[Tables 7-10](https://shorturl.at/JZufO)**).  \nThe findings indicate that (1) DPE is robust to both τ and α on Waterbirds and MetaShift, with low standard deviations across tested values; (2) Living17 shows greater sensitivity, likely due to its more challenging subpopulation structure. Nonetheless, even in this case, worst-group accuracy varies within an acceptable range.  \n**These results confirm that DPE’s performance is generally stable across a reasonable range of hyperparameter settings**. We’ll include this ablation study in the final version of the paper.\n\n### 3. Prototype-subgroup alignment\nWe appreciate the reviewer’s desire for deeper theoretical insight. To better understand whether learned prototypes capture meaningful subgroup structure, we conducted an exploratory qualitative analysis. Specifically, we tasked a third party (ChatGPT) with identifying common traits among the closest samples to each prototype—without providing group labels. This process revealed recurring semantic and ecological patterns (e.g., habitat type, pose, morphology) within prototype clusters, despite not being explicitly supervised to discover such groupings (see **[Figure 1 and 2](https://shorturl.at/o6tUI)**). These findings suggest that Diversified Prototypical Ensembles (DPE) may encourage meaningful prototype-subgroup alignment, potentially contributing to improved worst-group performance. We now clarify this insight in the revised draft.\n\n---\n\n**Thank you again for your time and thoughtful review. If our responses addressed your concerns, we’d appreciate your consideration in raising the score. Please don’t hesitate to let us know if any points remain unclear—we’re happy to provide further clarification.**", "claims_and_evidence": "The main claim is that DPE improves robustness to subpopulation shifts, as measured by WGA, compared to existing methods. The evidence provided is primarily empirical, based on experiments on several benchmark datasets. While the results show consistent improvements, the evidence is not entirely convincing due to concerns about the fairness of the comparisons (see below). The claims regarding the importance of diversification are supported by ablation studies.", "methods_and_evaluation_criteria": "The proposed method (DPE) is a reasonable approach for the problem. Combining prototypical networks with ensemble learning and explicit diversification is a novel and potentially effective strategy. The evaluation criteria, primarily worst-group accuracy (WGA), is appropriate for assessing robustness to subpopulation shifts. The choice of benchmark datasets (from Yang et al., 2023) is also standard in this area.", "theoretical_claims": "The paper does not present any formal proofs or theoretical claims.", "experimental_designs_or_analyses": "The authors state that their initial ERM training stage (\"ERM*\") uses stronger data augmentation and longer training than the ERM implementation in Yang et al. (2023). They use published results for several baselines (ERM, CRT, ReWeightCRT, DFR from Yang et al., 2023; RWY and AFR from their original papers), but it seems not all of them use this stronger protocol. This introduces a confounder, making it unclear whether improvements are due to DPE itself or the stronger ERM training.", "relation_to_broader_scientific_literature": "The paper builds upon several areas of related work, including subpopulation shift, prototypical networks, and ensemble learning.", "essential_references_not_discussed": "The paper appears to cover the most essential references for the core problem and proposed method.", "other_strengths_and_weaknesses": "1. The method introduces several hyperparameters: the number of prototypes (N), the temperature ($\\tau$), the IPS loss weight ($\\alpha$), and the size of the subsets used for training. While the authors provide some details on how these were chosen, the sensitivity to these parameters (except N) isn't thoroughly explored.\n2. The paper relies heavily on intuition and empirical results. There's no theoretical analysis explaining why the diversification strategies lead to improved worst-group accuracy, or if and why the prototypes align well with relevant subpopulations. A more formal understanding of the method's properties would be valuable.", "other_comments_or_suggestions": "None.", "questions_for_authors": "See above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AxhxJjD8bm", "reviewer": "Reviewer_aEu5", "summary": "This paper studies the subpopulation shifting problem. To alleviate the issue, motivated by the idea of ensemble learning, the author proposes using a mixture of diversified prototypical classifiers over the feature prototypes of the subpopulations to classify different subpopulations correctly. Extensive experiments have been conducted on standard datasets. The author provides a comprehensive comparison of the proposed methods over the state-of-the-art method and justifies the effectiveness of the proposed method.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive and constructive review. We appreciate your recognition of our well-motivated method, the novel application of prototype-based ensemble for subpopulation shift, and the thorough experimental validation.\nYour main concerns include \n- (1) lack of ablation on the number of prototypes ($N$)\n- (2) missing analysis of computational and storage cost\n- (3) absence of discussion on related work in incremental subpopulation learning and boosting-based prototype methods.\n\nWe address each of these points below with new results, expanded discussion, and clarifications. Full tables are provided at https://github.com/anonymous102030411/anon.\n\n### 1. Choosing the Number of Prototypes\nWe thank the reviewer for pointing out the importance of prototype count in achieving effective diversification. This is indeed a key design parameter in our method, and we have now conducted an extended ablation study to better quantify its effect on model performance. For a discussion of this study, we direct the reader to our response to reviewer ___[rN4y](https://shorturl.at/LBOK6)___.\n\n### 2. Computation Cost and Complexity\nWe agree that evaluating memory and computational requirements is essential. The additional cost of DPE relative to baselines is minimal, as our method only introduces $N$ prototypes per class ($D$-dimensional). In our experiments, with $K=10$ to $K=100$, $D=1024,$ and $N=10$ to $20$, this overhead is negligible compared to the backbone encoder.\nInference speed is largely unaffected, as matching features against prototypes adds minimal cost relative to the encoder’s forward pass. While prototype storage scales with $K$, our experiments confirm that memory and compute demands remain manageable in practical settings. We refer the reviewer to our response to ___[rN4y](https://shorturl.at/LBOK6)___, where we provide detailed benchmarking results, including inference speed and GPU memory usage (see **[Table 1](https://shorturl.at/KekLM)**). We will update the manuscript to clarify that computation cost remains low and include the full benchmarking table in the appendix.\n\n### 3. Addition of Relevant Citations\nWe appreciate the reviewer’s suggestion and acknowledge that [1] addresses subpopulation shifts in an incremental learning setting. It employs an ERM-trained feature extractor and incrementally updates the classifier by combining previous and newly trained classifiers, with margin-enforce loss focusing on hard examples. The update balance is determined by measuring prototype distortion under the new classifier.\nOur work differs from [1] as follows:\n- Subpopulation characteristics: While [1] assumes distinct, predefined subpopulations, our method operates without prior knowledge of subpopulation structures and successfully generalizes to unknown shifts (as shown in Table 1, main paper).\n - Training setup: We focus on distribution shift robustness within a single training set, unlike [1], which assumes incremental subpopulation learning.\n - Methodology: Instead of using prototypes to balance learning vs. forgetting, we employ prototype-based classification with explicit diversification to enhance subgroup discovery.\nIn response to the reviewer’s feedback, we will cite [1] and supporting references in the revised manuscript, highlighting how prototype-based methods and margin-enforce loss from boosting theory refine decision boundaries for subpopulation shift challenges, while clearly delineating our contributions. We hope this clarification strengthens the contextual positioning of our work.\n\n### Relation to Broader Scientific Literature\nThe reviewer highlights the relevance of our work to addressing the subpopulation shift problem. We would like to underscore the generality of the subpopulation shift phenomenon: as noted in [3], challenges such as class imbalance, attribute imbalance, and spurious correlations are all instances of the more general phenomenon of subpopulation shift. By targeting these core issues, our approach has broad applicability across domains like medical imaging, fairness in hiring, and large-scale recognition tasks, ultimately leading to more reliable, inclusive models.\n\n[3] Yang et al., 2023. Change is hard: a closer look at subpopulation shift\n\n---\n\n**Thank you for taking the time to review our work. If our responses resolve your concerns, please kindly consider raising your score. Please feel free to reach out with any remaining questions.**", "claims_and_evidence": "The paper's central claim is that explicitly encouraging diversity of the prototype-based predictors for each class could encourage subsequent ensemble members to capture the different decision boundaries for each subgroup, leading to better performance on classification under subpopulation shifting. The claim is reasonable, and the author provides extensive empirical studies to justify its effectiveness. Figure 2 provides a concise and clear demonstration of the motivation of the claim and the proposed method. The ablation study and error analysis further demonstrate that the proposed methods improve the classification performance as expected.", "methods_and_evaluation_criteria": "Strengths:\n\nOverall, the paper is well-written and well-motivated. The proposed prototype-based method is relevant and sound, and the idea of ensemble diversified prototype-based predictors to alleviate the subpopulation shifting issue is originated from the ensemble learning perspective but with further consideration of the subpopulation shifting problem. Extensive experiments have been conducted to verify the effectiveness of the proposed methods from different perspectives.\n\n\nWeaknesses:\n\n1. In line 220, the author chooses `N` prototypes per class for `K` classes. However, the author does not elaborate on how the `N` is being chosen, and the impact of the choice of `N` for training performance is unclear. The author may want to provide further ablation studies to demonstrate it, as the choice of `N` should be highly relevant to whether we can exactly achieve diversified predictors in practice, e.g., when `N` is not large enough, it is hard to achieve diversified predictors as there is not sufficient prototype to represent the feature space of the subgroups.\n\n2. The cost of the proposed method has not been discussed. As the proposed method is prototype-based, the author must save the prototype and predictors to perform inference in test time. Given that the compared methods are not prototype-based, the author needs to provide the computational and storage costs between the proposed methods and the compared methods to let the reader know the trade-off between performance and the storage and computation cost among different kinds of methods to demonstrate the practicality of the proposed method.\n\n3. The missing relevant literature in subpopulation shifting. Please refer to the Essential References Not Discussed for details.", "theoretical_claims": "This is not a theory paper, and there is no theoretical analysis.", "experimental_designs_or_analyses": "The experimental designs and analysis are sound.", "relation_to_broader_scientific_literature": "The submission is specific to addressing the subpopulation shifting problem and may not significantly impact broader scientific literature.", "essential_references_not_discussed": "Paper [1] has not been cited/discussed in the paper. [1] tackles the subpopulation shifting problem in an incremental learning manner. In [1], the author proposes a prototype-based incremental learning method based on the generalized boosting theory [2,3] to learn new classifiers for novel subpopulations incrementally and combine the old and new classifiers sequentially. The insight of [1] is to calibrate the decision boundary over old and new subpopulations to correctly classify different subpopulations. Such a fine-grain modification on decision boundary is achieved by optimizing the margin-enforce loss [2,3], which, in theory, is equivalent to performing the ensemble learning via gradient-boosting that minimizes the residual error of the previously learned classifier.\n\nThe reviewer recognizes that the present submission is not a continual or incremental learning paper. However, considering that both [1] and the present submission have the same research goal of addressing subpopulation shifting and similar insight on alleviating the subpopulation shift issue by ensemble learning and prototype-based methods, the reviewer believes the author must provide sufficient discussion of the current submission and [1] to properly acknowledge the existing literature and inspire the boarder research community.\n\n[1] Balancing between forgetting and acquisition in incremental subpopulation learning. ECCV 2022\n\n[2] Multiclass boosting: margins, codewords, losses, and algorithms. JMLR 2019\n\n[3] Boosting: Foundations and Algorithms. Kybernetes 2013", "other_strengths_and_weaknesses": "N/A.", "other_comments_or_suggestions": "N/A.", "questions_for_authors": "N/A.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vkfpBL6bTQ", "reviewer": "Reviewer_mdbr", "summary": "The paper introduces **Diversified Prototypical Ensemble (DPE)** to enhance robustness against subpopulation shifts. The method trains multiple diverse prototypes per class on top of a frozen feature extractor and enforces feature diversity through **inter-prototype similarity (IPS) loss**. By restructuring the feature space, DPE ensures that minority subgroups receive dedicated representations, improving worst-group accuracy (WGA). Experiments across datasets like **Waterbirds and CelebA** show that DPE **outperforms prior reweighting-based approaches** in handling subpopulation imbalances. **Ablation studies confirm that prototype diversification and ensemble strategies drive these gains.** While effective, the authors highlight limitations in **prototype selection strategies and computational efficiency**, which require further optimization across diverse datasets.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed and constructive feedback on our submission. We appreciate the time you have dedicated to evaluating our work, and we are pleased that you recognize the strength of our method in improving worst-group accuracy (WGA) under subpopulation shifts. Your concerns are \n- (1) Unclear contribution of DPE versus augmentations in ERM*\n- (2) Need for ablations on the number of prototypes\n- (3) Missing evaluation on other BREEDS datasets\n- (4) Absence of standard accuracy reporting\n- (5) A missing reference to prior work involving hierarchical learning\n- (6) Lack of interpretability analysis (e.g., t-SNE). \n\nWe address these points below with new experiments, clarifications, and discussion. Full tables are provided at https://github.com/anonymous102030411/anon.\n\n### 1. ERM vs. ERM∗ (Backbone Confounding Issue):\nTo directly evaluate whether DPE’s performance gains stem from its architecture or stronger backbones, we have retrained the feature extractors for all datasets using the original ERM configurations from Yang et al. (2023)—without additional augmentations or extended training. The results using DPE on these standard ERM features are now included as ERM + DPE in the revised Table 2 and Table 3, (see [Tables 2-3](https://shorturl.at/Dy4Ab)). Across all datasets, DPE still achieves substantial gains over baselines under the same backbone. For example, on Waterbirds, ERM yields 69.1±4.7 WGA, while ERM + DPE achieves 91.0±0.5. Similarly, on CivilComments, DPE improves WGA from 63.2±1.2 (ERM) to 71.5±0.6. **These findings confirm that the improvement comes from DPE itself, not just stronger pretraining**. The updated results are detailed in our response to reviewer [VrPG](https://openreview.net/forum?id=qUTiOeM57J&noteId=JePMn0uimT).\n\n### 2. Prototype Count ($N$) Ablation\nFollowing the reviewer’s suggestion, we have expanded our ablation on the number of prototypes N, which will be included in the appendix. For a more detailed discussion of our analysis, we direct the reader to the response to reviewer ___rN4y___.\n\n### 3. Additional BREEDS Evaluation – Non-living 26 and Entity-30\nWe thank the reviewer for this suggestion. BREEDS benchmarks such as Living-17, Non-living-26, and Entity-30 are important for studying subpopulation shifts as they construct domain splits based on fine-grained taxonomies, inducing subgroup variation. To address the comment, we conducted new experiments on Non-living-26 and Entity-30 using ERM and ERM + DPE, with and without ImageNet pretraining. Results show that DPE consistently improves worst-group accuracy (e.g., +3%–10%) across both source and target domains (see [Tables 4-5](https://shorturl.at/yuYnY)), These results validate DPE's robustness and generalization under complex subpopulation shift settings and will be included in the final version of the paper.\n\n### 4. Reporting Standard Accuracy\nWe thank the reviewer for highlighting the importance of reporting standard (average) accuracy to better contextualize the gains in worst-group accuracy (WGA). In response, we have included average accuracy results across all benchmarks in the appendix of the revised version. The average accuracy table corresponding to Table 2 is shown in [Table 6](https://shorturl.at/4CFih). The results demonstrate that both ERM + DPE and ERM* + DPE maintains comparable or improved average accuracy relative to ERM and ERM* across all datasets. These results confirm that the improvements in worst-group accuracy offered by DPE do not come at the cost of overall accuracy. On the contrary, DPE often improves both metrics. We include these results and observations in the revised manuscript.\n\n### 5. Hierarchical Learning Literature\nWe thank the reviewer for highlighting the connection to hierarchical representation methods. While DPE does not rely on explicit taxonomies, we agree it shares the goal of feature space structuring. We now cite and discuss \"Encoding Hierarchical Information in Neural Networks Helps in Subpopulation Shift\" in the Related Work section. We clarify that unlike those methods, DPE infers structure from data without pre-defined hierarchies, enabling application to a broader set of tasks lacking such annotations.\n\n### 6. Feature Interpretability via t-SNE and Prototype Visualization\nTo clarify the intuition behind the prototype ensemble’s ability to capture semantically meaningful and generalizable subpopulation features, we note that Figures 1.4 and 1.5 are real T-SNE visualizations from our Waterbirds experiments. They highlight prototypes aligned with semantic concepts (e.g., “small yellow bird”) rather than spurious ones (e.g., “land background”). We will include full-size T-SNE plots and additional examples from other datasets in the appendix.\n\n---\n\n**Thank you for taking the time to review our work. If our answers resolve your concerns, we’d appreciate your consideration in raising the score. We're happy to clarify any remaining questions.**", "claims_and_evidence": "1. DPE improves WGA compared to re-weighted approaches. Tables 1 and 2 support this claim. \n 2. Figure 5 demonstrates that IPS loss encourages diverse feature representations.", "methods_and_evaluation_criteria": "1. While DPE improves WGA, interpret ability analysis is lacking. If t-SNE or feature visualization plots be shared then it will add to the understanding. \n\n2. DPE builds on $ERM^*$ and not on ERM. The advanced data augmentations affects learned features in an ERM setting as shown in works such as CutMix. DPE’s effectiveness on standard ERM should be tested to separate augmentation effects from prototype diversification.\n\n3. How are N prototypes determined? An ablation on N across datasets would clarify its optimal selection.\n\n4. Only Living-17 is used as one of the BREEDs benchmarks. Other datasets such as Non-living 26 and Entity-30 should be addressed as well. These are subpopulation shift vision benchmarks. One can use datasets such as CELEBA as a benchmark, but evaluation on either on Non-lIving 26 or Entity 30 is required.", "theoretical_claims": "No theoretical claims made.", "experimental_designs_or_analyses": "The experimental section overall confirms DPE improves WGA against previous approaches. \n\n1. For gap in analysis, please refer to the experiments related to Non-living 26 or Entity 30 and feature visualizations via t-SNE. \n2. A major issue is the effect of $ERM^(*)$ as feature extractor and not ERM. Experimental validation is required here. \n3. Please provide standard accuracy as well. Without that, it is difficult to understand the drop for WGA. \n\nThese experiments are missing and will help in understanding the process.", "relation_to_broader_scientific_literature": "Very relevant for domain shift literature. Assuming subpopulation shift falls as a sub problem under domain shift.", "essential_references_not_discussed": "Since prior work (e.g., hierarchical learning methods) has addressed subpopulation shift by structuring feature space, this paper should acknowledge that conceptual and ideally experimental overlap. While hierarchical classification relies on predefined label hierarchies, prototypical classifiers structure feature space through learned prototypes. Although implementation details differ, both approaches share the goal of improving generalization under subpopulation shifts, making the connection relevant. This point becomes as a novelty issue for this approach.\n\n1. \"Encoding Hierarchical Information in Neural Networks Helps in Subpopulation Shift\", IEEE Transactions on Artificial Intelligence, 2023/ Fine Grained Visual Categorization (FGVC9) workshop, CVPR 2022.", "other_strengths_and_weaknesses": "1. Figure 2 is very well illustrated and conveys the goal of the paper. \n2. Limitations are clearly explained.", "other_comments_or_suggestions": "1. Why is ERM and ERM* on LIVING-17 different from the original BREEDs paper? \n2. Please explain the numbers in Figure 3. In Table 2, gap between ERM* and DPE is aorund 10% for LIVING-17. Why is it showing 18+ on Figure 3? \n\n3. For Figure 5. what is the number of prototypes?", "questions_for_authors": "Please address the gaps in experiments provided under Experimental Designs and Analyses. \n1. Results on Non-living 26 \n2. Results of ERM plus PDE on any 2 datasets \n\nThis should help us understand the efficacy of approach.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper tackles the problem of subpopulation shifts by training a diverse ensemble of prototype classifiers on top of a pre-trained feature extractor. The core idea is to have each prototype focus on different aspects of the data, thereby discovering and correctly classifying diverse subpopulations without explicit annotations. The authors demonstrate this approach can achieve better worst-group accuracy compared to existing methods. 

The review team generally found the work to be well-motivated, with a promising methodology that combines prototypical classifiers with ensemble learning and explicit diversification strategies. The extensive empirical evaluations across nine diverse real-world datasets were highlighted as a strength. One key concern revolved around the lack of grounded justification or intuition for the proposed method and why the diversification strategies lead to improved worst-group accuracy. While I won’t repeat all dimensions of improvement called out in the review process, I urge the authors to work on additional ablation studies on the number of prototypes, HP sensitivity, and the potential confounding effect of the stronger ERM method.  As a general practice, include average accuracy alongside worst-group accuracy will provide additional context for the readers.
Decision: Accept (poster)