# ProME: Prototype-Margin Environments with Repair-Aware Selection for Group-Robust Learning

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 34 of its 48 references that are in the map.
- First posted on arXiv 13 Aug 2026.

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

- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [Prototypical Networks for Few-shot Learning](https://mlsanity.com/paper/1703.05175/) (2017)
- [Shortcut Learning in Deep Neural Networks](https://mlsanity.com/paper/2004.07780/) (2020)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [Prevalence of Neural Collapse during the terminal phase of deep learning training](https://mlsanity.com/paper/2008.08186/) (2020)
- [Causal inference using invariant prediction: identification and confidence intervals](https://mlsanity.com/paper/1501.01332/) (2015)
- [Just Train Twice: Improving Group Robustness without Training Group Information](https://mlsanity.com/paper/2107.09044/) (2021)
- [Nuanced Metrics for Measuring Unintended Bias with Real Data for Text Classification](https://mlsanity.com/paper/1903.04561/) (2019)
- [Last Layer Re-Training is Sufficient for Robustness to Spurious Correlations](https://mlsanity.com/paper/2204.02937/) (2022)
- [Environment Inference for Invariant Learning](https://mlsanity.com/paper/2010.07249/) (2020)
- [The Risks of Invariant Risk Minimization](https://mlsanity.com/paper/2010.05761/) (2020)
- [No Subclass Left Behind: Fine-Grained Robustness in Coarse-Grained Classification Problems](https://mlsanity.com/paper/2011.12945/) (2020)
- [Simple data balancing achieves competitive worst-group-accuracy](https://mlsanity.com/paper/2110.14503/) (2021)
- [Correct-N-Contrast: A Contrastive Approach for Improving Robustness to Spurious Correlations](https://mlsanity.com/paper/2203.01517/) (2022)
- [Does Invariant Risk Minimization Capture Invariance?](https://mlsanity.com/paper/2101.01134/) (2021)
- [Spread Spurious Attribute: Improving Worst-group Accuracy with Spurious Attribute Estimation](https://mlsanity.com/paper/2204.02070/) (2022)
- [Simple and Fast Group Robustness by Automatic Feature Reweighting](https://mlsanity.com/paper/2306.11074/) (2023)
- [ZIN: When and How to Learn Invariance Without Environment Partition?](https://mlsanity.com/paper/2203.05818/) (2022)
- [Towards Last-layer Retraining for Group Robustness with Fewer Annotations](https://mlsanity.com/paper/2309.08534/) (2023)
- [Model Agnostic Sample Reweighting for Out-of-Distribution Learning](https://mlsanity.com/paper/2301.09819/) (2023)
- [The Clever Hans Mirage: A Comprehensive Survey on Spurious Correlations in Machine Learning](https://mlsanity.com/paper/2402.12715/) (2024)
- [Discovering environments with XRM](https://mlsanity.com/paper/2309.16748/) (2023)
- [Group Robust Classification Without Any Group Information](https://mlsanity.com/paper/2310.18555/) (2023)
- [Improving Group Robustness on Spurious Correlation Requires Preciser Group Inference](https://mlsanity.com/paper/2404.13815/) (2024)
- [Learning Robust Classifiers with Self-Guided Spurious Correlation Mitigation](https://mlsanity.com/paper/2405.03649/) (2024)
- [Out of Spuriousity: Improving Robustness to Spurious Correlations without Group Annotations](https://mlsanity.com/paper/2407.14974/) (2024)
- [Robustness to Subpopulation Shift with Domain Label Noise via Regularized Annotation of Domains](https://mlsanity.com/paper/2402.11039/) (2024)
- [Bias Amplification Enhances Minority Group Performance](https://mlsanity.com/paper/2309.06717/) (2023)
- [Group-robust Sample Reweighting for Subpopulation Shifts via Influence Functions](https://mlsanity.com/paper/2503.07315/) (2025)

## Abstract (arXiv)

Group-robust learning is crucial for maintaining accuracy on rare subpopulations when training-group labels are unavailable. However, existing methods often infer environments from a separate reference model and select representations before fitting the classifier used at deployment, leaving both decisions misaligned with the deployed predictor. In this work, we formulate group robustness without training-group labels as the endogenous environments with repair-aware selection (ERAS) problem, and propose ProME (Prototype-Margin Environments) to align both decisions with the deployed predictor. ProME splits prototype margins at their median to construct approximately balanced environments along the training trajectory, and fits a group-balanced linear head on group-annotated validation data to rank the resulting predictors by validation worst-group accuracy. We theoretically bound the worst risk across the inferred environments for a fixed predictor and partition, showing that this bound transfers to the oracle groups under an explicit alignment condition. Extensive experiments show that prototype margins enrich shortcut-conflicting examples, classifier repair reshapes candidate evaluation, and ProME achieves the highest average worst-group accuracy among the compared methods with the same group-label access.