Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift

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

• (1) Unclear contribution of DPE versus augmentations in ERM*
• (2) Need for ablations on the number of prototypes
• (3) Missing evaluation on other BREEDS datasets
• (4) Absence of standard accuracy reporting
• (5) A missing reference to prior work involving hierarchical learning
• (6) Lack of interpretability analysis (e.g., t-SNE).

We address these points below with new experiments, clarifications, and discussion. Full tables are provided at https://github.com/anonymous102030411/anon.

1. ERM vs. ERM∗ (Backbone Confounding Issue):

To 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). 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.

2. Prototype Count ($N$) Ablation

Following 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.

3. Additional BREEDS Evaluation – Non-living 26 and Entity-30

We 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), These results validate DPE's robustness and generalization under complex subpopulation shift settings and will be included in the final version of the paper.

4. Reporting Standard Accuracy

We 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. 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.

5. Hierarchical Learning Literature

We 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.

6. Feature Interpretability via t-SNE and Prototype Visualization

To 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.

---

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.

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. Your main concerns include

• (1) lack of ablation on the number of prototypes ($N$)
• (2) missing analysis of computational and storage cost
• (3) absence of discussion on related work in incremental subpopulation learning and boosting-based prototype methods.

We address each of these points below with new results, expanded discussion, and clarifications. Full tables are provided at https://github.com/anonymous102030411/anon.

1. Choosing the Number of Prototypes

We 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.

2. Computation Cost and Complexity

We 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. Inference 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, where we provide detailed benchmarking results, including inference speed and GPU memory usage (see Table 1). We will update the manuscript to clarify that computation cost remains low and include the full benchmarking table in the appendix.

3. Addition of Relevant Citations

We 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. Our work differs from [1] as follows:

• 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).
• Training setup: We focus on distribution shift robustness within a single training set, unlike [1], which assumes incremental subpopulation learning.
• Methodology: Instead of using prototypes to balance learning vs. forgetting, we employ prototype-based classification with explicit diversification to enhance subgroup discovery. In 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.

Relation to Broader Scientific Literature

The 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.

[3] Yang et al., 2023. Change is hard: a closer look at subpopulation shift

---

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.

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.

Your main concerns include

• (1) potential confounding due to stronger ERM* training compared to baseline implementations
• (2) limited hyperparameter sensitivity analysis, particularly for τ, α, and subset size
• (3) unclear prototype-subgroup alignment.

We address each of these points below with new results, analyses, and clarifications. Full tables are provided at https://github.com/anonymous102030411/anon.

1. Ensuring a Fair Comparison with Prior Art

We 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) confirms that DPE outperforms all reported baselines, showing the gains stem from DPE itself, not the ERM backbone.

We also report results for ERM* + DPE, which includes stronger augmentation and longer training (see Table 2b). While these features boost performance, DPE’s advantage remains consistent across both ERM and ERM* setups.

Regarding fairness in comparing ERM* + DPE to prior work, we note that several baselines use enhanced protocols:

• DFR includes data augmentation.
• RWG and RWY use extended training and thorough tuning.
• CRT’s code includes augmentation, though not mentioned in the paper.
• CnC, SSA, and GAP lack public code, so augmentation use is unclear.

Thus, 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. To support transparency, we will include:

• A clear discussion of ERM vs. ERM* in the paper.
• A full performance table with both ERM + DPE and ERM* + DPE results.

2. Hyperparameter Tuning and Sensitivity

We 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.
To 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).
The 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.
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.

3. Prototype-subgroup alignment

We 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). 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.

---

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.

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.
You raised key concerns regarding:

• (1) ablation on prototype count and hyperparameter sensitivity,
• (2) missing runtime/resource analysis,
• (3) incomplete reporting of average accuracy,
• (4) limited evidence linking prototypes to meaningful subpopulations, and
• (5) lack of theoretical justification.

We address each of these below with new results, analyses, and clarifications. Full tables are provided at https://github.com/anonymous102030411/anon.

1. Further Ablation on the Number of Prototypes

We 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.
Let $\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:

$\Delta_N = \frac{\text{WGA}_N - \text{WGA}_1}{\text{WGA}_1} \times 100\%$

We 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:

$\Delta_5 = 2.4\%$, $\Delta_{10} = 3.3\%$, $\Delta_{15} = 3.7\%$, $\Delta_{25} = 3.7\%$, $\Delta_{40} = 3.7\%$

These 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.

We 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. Regarding the performance sensitivity to other important hyperparameters in our study, we direct the reviewer to our response to Reviewer VrPG with more extensive results and discussion.

2. Computational Complexity Analysis

We 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.

Table 1 (See Table 1) 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.

3. Reporting Average Accuracy

We 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), demonstrating that DPE remains competitive on average accuracy while prioritizing worst-group robustness.

4. Clarifying the Link Between Learned Prototypes and Subpopulations

Our 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). This supports our central claim that diversity in prototype space enables better subgroup coverage. Expanded ablations further support the empirical foundation of our method.

5. Theoretical Justification

The 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.

---

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.