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

ProME: Prototype-Margin Environments with Repair-Aware Selection for Group-Robust Learning Qianqian Wanga,b , Yunshan Lib , Dawei Huangb , Wenwu Gongb,∗ and Lili Yanga,b,∗ a Shenzhen Key Laboratory of Safety and Security for Next Generation of Industrial Internet, Shenzhen, China

ARTICLE INFO

ABSTRACT

Keywords: group robustness spurious correlations invariant risk minimization repair-aware model selection prototype margin out-of-distribution generalization

Group-robust learning is crucial for maintaining accuracy on rare subpopulations when traininggroup 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.

Group-robust learning seeks to preserve accuracy across label–attribute groups whose proportions differ between training and deployment, as in medical diagnosis, content moderation, and attribute recognition [1, 2]. Rare groups contribute little to the average loss, allowing empirical risk minimization (ERM) to achieve high overall accuracy while accuracy on rare groups remains low. This disparity often arises when ERM relies on shortcut features whose correlations with labels hold for common groups but fail for rare groups [3]. For example, background and gender can serve as shortcut features for predicting bird type and hair color in Waterbirds and CelebA, respectively [1, 4, 5]. Errors caused by these shortcuts have greater impact when rare groups become more prevalent at deployment. Worst-group accuracy (WGA) captures this vulnerability by reporting the accuracy of the least accurate group instead of averaging across groups [1]. Fig. 1 illustrates this setting through the Waterbirds training-group proportions and reported WGA on three benchmarks. The challenge is to maintain WGA as group proportions shift. Group distributionally robust optimization (Group DRO) targets WGA by minimizing the worst empirical group risk [1, 6, 7]. It requires a group label for every training example, yet such labels are costly or infeasible to collect for large training datasets. We use oracle groups to denote the label–attribute groups whose labels are unavailable during training. Methods without training-group labels commonly ∗ Corresponding authors

12331103@mail.sustech.edu.cn (Q. Wang); 12331105@mail.sustech.edu.cn (Y. Li); 12312405@mail.sustech.edu.cn (D. Huang); gongww@sustech.edu.cn (W. Gong); yangll@sustech.edu.cn (L.

Yang)

ORCID (s):

Q. Wang et al.: Preprint submitted to Elsevier

use a two-stage pipeline. A reference model infers environments by clustering learned features [8–10] or identifying misclassified examples [6, 7, 11, 12]. A separate predictor is trained using the resulting fixed assignments [6, 8, 9]. This separation creates an environment–representation mismatch. This mismatch can weaken invariant learning [13, 14] and limit WGA on the oracle groups. (a)

80 73.0

LW

60 40 22.0

20 0

WL

3.8 1.2

LL LW WL WW

LL: landbird on land LW: landbird on water

WL: waterbird on land WW: waterbird on water

worst-group accuracy (%)

1. Introduction

group proportion (%)

arXiv:2608.13190v1 [cs.LG] 13 Aug 2026

b Department of Statistics and Data Science, Southern University of Science and Technology, Shenzhen, China

100

(b)

80

60

40

CivilComments

CelebA

ERM GroupDRO

Waterbirds

DFR ProME (ours)

Fig. 1. Training-group imbalance and worst-group accuracy. (a) Group proportions in the Waterbirds training set. (b) Reported WGA of representative methods, including deep feature reweighting (DFR), on three benchmarks.

A second misalignment arises during model selection, when an encoder checkpoint must be chosen for deployment. An encoder checkpoint with low pre-repair WGA can still retain features that support a stronger classifier after repair. Classifier repair refits the last layer on group-balanced data [15–17]. Selecting encoder checkpoints by pre-repair WGA, or retaining the final checkpoint by default, can reject the checkpoint with the highest post-repair WGA. Together, the environment–representation mismatch and prerepair model selection leave both decisions misaligned with the deployed predictor. This limitation raises a pivotal issue: Page 1 of 17

ProME: Group Robustness without Training Group Labels

How can we align environment construction and model selection with the deployed predictor without traininggroup labels? In this work, we formulate group robustness without training-group labels as the endogenous environments with repair-aware selection (ERAS) problem to align both decisions with the deployed predictor (Definition 1), and propose ProME (Prototype-Margin Environments) to implement this formulation. In ERAS, an environment partition is endogenous when an encoder checkpoint creates the partition used to regularize later checkpoints on the same trajectory. Model selection is repair-aware when encoder checkpoints are ranked by WGA after classifier repair. To construct endogenous environments, ProME assigns every example a prototype margin between its observed class and the strongest competing class under a cosine-prototype classifier. A median split of these margins produces two approximately balanced environments for invariant learning. In Stage 2, ProME freezes each retained encoder checkpoint, fits a group-balanced linear head on group-annotated validation data, and deploys the encoder–head pair with the highest validation WGA (Eq. (18)). Environment construction and representation learning use only inputs and class labels. Group-annotated validation data support candidate tracking, classifier repair, regularization tuning, and selection; they do not update the Stage 1 representation. Overall, ProME constructs environments without traininggroup labels, ranks candidates after classifier repair, and deploys a single encoder–head pair. Theoretically, we characterize three Stage 1 properties under explicit conditions. First, low prototype margins mark shortcut-conflicting examples when the causal and spurious components are sufficiently distinguishable (Proposition 1). Second, the induced assignments change by a controlled amount as the representation changes (Proposition 2). Third, for a fixed predictor and partition, we bound the worst inferred-environment risk by the mean and variance of the environment risks (Proposition 3). Under an explicit totalvariation alignment condition, this bound transfers to the oracle groups (Corollary 1). Classifier repair and repairaware selection in Stage 2 are evaluated experimentally. To evaluate ProME and examine its two alignment mechanisms, we study three real-world benchmarks (Waterbirds [1], CelebA [5], and CivilComments [18]) and one controlled benchmark (ColoredMNIST [19]). The main comparison uses published baseline results under their original protocols. Controlled studies use matched settings. They show that trajectory-derived margins improve pre-repair WGA in the matched source comparison and that classifier repair reshapes candidate evaluation. Among the methods compared with the same group-label access, ProME raises average WGA from 83.9% to 87.0%. Our contributions are summarized as follows:

ERAS aligns environment construction and model selection with the deployed predictor. • Two-stage ProME framework. We propose ProME, a two-stage framework that constructs approximately balanced environments from prototype margins and ranks encoder–head pairs after classifier repair. It requires no training-group labels, is backbone-agnostic, and deploys only the selected encoder–head pair. • Mechanism-aligned theory. We establish theoretical results. Under explicit conditions, low prototype margins mark shortcut-conflicting examples, and the induced partition changes by a controlled amount with the representation. We also bound the worst inferredenvironment risk for a fixed predictor and partition. This bound transfers to oracle groups under an explicit alignment condition.

2. Related Work

ProME builds on three lines of work: environment inference without training-group labels, prototype geometry for environment inference, and classifier repair with groupannotated validation data. Environment inference without training-group labels. When group labels are available, Group DRO directly optimizes the worst empirical group risk, and group-balanced training provides a strong supervised baseline [1, 20]. Without traininggroup labels, error-based methods identify hard or misclassified examples with a reference classifier and use them for reweighting or contrastive learning [6, 11, 21, 22]. Other methods exploit feature clusters, confidence, disagreement, or pretrained representations to identify spurious structure or estimate pseudo-groups [8, 10, 23–26]. Counterfactual alignment provides a complementary diagnostic for spurious correlations in black-box classifiers [27]. Environmentinference methods instead optimize a partition through invariant-learning or cross-risk signals [7, 9, 28, 29]. In these pipelines, an auxiliary model, reference predictor, or clustering stage produces environments that are subsequently used by a robust objective. This separation can create an environment–representation mismatch when the partition does not track the representation being optimized. ProME derives its inferred environments from the representationlearning trajectory itself. prototype geometry for environment inference. Invariant learning seeks predictors that remain stable across environments [19, 30–33]. Learning such predictors without a supplied environment partition requires additional information or an explicit inductive bias [34], and weak variation in the relevant shortcut can leave invariant objectives with a nonrobust solution [35, 36]. Prototypical classifiers summarize each class by a mean representation, while neural collapse • Deployment-aligned formulation. We formulate group relates class means to terminal classifier geometry [37, 38]. Diverse prototypical ensembles further use multiple class robustness without training-group labels as endogenous environments with repair-aware selection (ERAS). prototypes to capture latent subpopulations [39]. ProME gives prototype geometry a different role: the observed-class

Q. Wang et al.: Preprint submitted to Elsevier

Page 2 of 17

ProME: Group Robustness without Training Group Labels

prototype margin from the current representation-learning trajectory defines the environments used for invariant learning. Classifier repair and repair-aware selection. ERM representations can retain core features even when the corresponding classifier has poor worst-group accuracy [15, 40]. Deep feature reweighting (DFR) freezes the encoder and retrains its last layer on group-balanced held-out data [15]. Subsequent work reduces the annotation requirements of last-layer retraining or revisits its implementation choices [16, 41, 42]. Automatic feature reweighting (AFR), group-aware priors, and related two-stage corrections also improve group robustness through a repaired or reweighted final layer [12, 43, 44]. These studies establish the value of classifier repair, yet the last-layer methods above typically evaluate a fixed or preselected representation. Their setup leaves model selection tied to a fixed or preselected representation. ProME fits the same type of group-balanced linear head to every retained encoder and ranks the resulting encoder–head pairs by validation WGA. This repair-aware comparison aligns model selection with the deployed predictor. Stage 1 environment construction and representation learning use no training-group labels; group-annotated validation data support candidate tracking, classifier repair, regularization tuning, and selection without updating the Stage 1 representation.

3. Preliminaries In this section, we introduce subpopulation shift and group-robust risk, review invariant learning when environments are observed, and define endogenous environment inference and repair-aware selection when they are not. For reference, The notations used in this paper are presented in Table 1.

3.1. Subpopulation Shift Let tr = {(𝑥𝑖 , 𝑦𝑖 )}𝑛𝑖=1 be the observed training set, where 𝑥𝑖 ∈  and 𝑦𝑖 ∈ {0, 1}. Each benchmark defines a finite collection  of oracle evaluation groups from the target and one or more spurious attributes. Stage 1 observes only (𝑥𝑖 , 𝑦𝑖 ). Group membership is available in the validation set for model selection and in the test set for evaluation. Let 𝑄𝑔 denote the conditional distribution of group 𝑔. When  partitions the training population, let 𝜋𝑔tr denote its training proportion. The training distribution is then the mixture ∑ ℙtr = 𝜋𝑔tr 𝑄𝑔 . (1) 𝑔∈

Under subpopulation shift, the group proportions may change at deployment even when the group-conditional distributions remain fixed. A predictor that performs well on the training mixture can still fail on a rare group whose correlation structure differs from that of the majority [1]. Average training accuracy hides these failures, and the weakest group decides whether the predictor we finally deploy is usable. Group robustness. For any distribution index 𝑞, let 𝑃𝑞 denote the corresponding distribution over (𝑋, 𝑌 ). Given a Q. Wang et al.: Preprint submitted to Elsevier

representation 𝜙, a linear classifier 𝑤, and a surrogate loss 𝓁, define [ ] 𝑞 (𝜙, 𝑤) ∶= 𝔼(𝑋,𝑌 )∼𝑃𝑞 𝓁(⟨𝑤, 𝜙(𝑋)⟩, 𝑌 ) . (2) ERM minimizes the average training risk. Group-robust learning instead controls the largest group risk: [ ] tr (𝜙, 𝑤) ∶= 𝔼ℙtr 𝓁(⟨𝑤, 𝜙(𝑋)⟩, 𝑌 ) , (3) wg (𝜙, 𝑤) ∶= max 𝑔 (𝜙, 𝑤). 𝑔∈

The accuracy counterpart of wg is worst-group accuracy WGA(ℎ) ∶= min

Pr

𝑔∈ (𝑋,𝑌 )∼𝑄𝑔

[ℎ(𝑋) = 𝑌 ] ,

where ℎ is the hard classifier induced by the prediction score. Both wg and WGA need oracle group identities, which Stage 1 does not have. They remain available for validation and evaluation. Without training-group labels, the learning problem shifts from optimizing known groups to building training environments that can stand in for them. The only signal available for that construction is the geometry of the representation being trained.

3.2. Problem Formulation Standard invariant-learning methods assume access to a partition of the training data into environments tr across which predictive correlations vary. Given these assignments, invariant risk minimization (IRM) [19] learns a representation 𝜙 and classifier 𝑤 using the scale-gradient objective ∑ 𝑒 (𝜙, 𝑤) IRM = 𝑒∈tr

+𝜆

∑‖ | ‖2 ‖∇𝛼 𝑒 (𝜙, 𝛼𝑤)| ‖ , ‖ |𝛼=1 ‖2

(4)

𝑒∈tr

where 𝛼 is the scalar dummy classifier and 𝜆 ≥ 0 controls the IRMv1 penalty. While well defined when environment assignments are observed, Eq. (4) does not specify how to construct them. To supply the missing assignments, a natural strategy is to read them off a detached reference encoder 𝜙ref trained beforehand. However, a partition fixed by 𝜙ref reflects the geometry of that encoder, not that of the representation trained downstream. This environment– representation mismatch then persists for the rest of training. We call such a construction exogenous, and reserve endogenous for a partition drawn from a checkpoint on the trajectory that it later regularizes. To remove the mismatch, we make environment construction part of the representation-learning problem rather than a preprocessing step. Let Π map a representation to a two-cell partition of tr , and let inv (𝜙, ) denote an invariance penalty evaluated on environments . For compactness, tr (𝜙) denotes the training risk of the prediction rule associated with 𝜙. Let  = {𝜃𝑡 } denote the retained encoder candidates, let 𝑓𝑡 be the Stage 1 classifier attached to 𝜃𝑡 , and let ̂ ℎ𝑡 be a group-balanced head fitted after freezing 𝜃𝑡 . We formalize the resulting problem as follows. Page 3 of 17

ProME: Group Robustness without Training Group Labels Table 1 Notations. Notation

Explanation

Notation

𝑥∈ 𝑎 tr , val

input spurious attribute used to define oracle groups training data without group labels and validation data with group labels backbone and projection head fixed auxiliary encoder used by exogenous environment-inference methods temperature of the prototype classifier normalized prototype of class 𝑐 split into low- and high-margin cells at 𝑚, Eq. (12) invariance-penalty weight at training step 𝑡 average training risk and worst-group risk

𝑦 ∈ {0, 1} binary target label 𝑔∈ oracle group;  is the set of oracle groups distribution of oracle group 𝑔 and its training pro𝑄𝑔 , 𝜋𝑔tr portion 𝜙𝜉,𝜔 normalized representation map defined by 𝑏𝜉 and 𝑝𝜔 𝑠(𝑥; 𝑦) observed-class prototype similarity minus the strongest competing similarity, Eq. (11) 𝑓𝜙 (𝑥) prototype logit, Eq. (10) 𝑚 median prototype margin on tr ̂0 , ̂1 inferred low- and high-margin environment cells 𝑃𝑞 , 𝑞 distribution indexed by 𝑞 and its expected loss IRM , REx IRM gradient penalty and risk extrapolation (REx) variance penalty 𝜌 upper bound on the total-variation distance between an oracle group and a matched inferred cell  set of candidate encoders retained from Stage 1 𝗆𝗈𝖽𝖾 environment-update policy used in Algorithm 1 WGAval validation WGA used to rank repaired encoder–head pairs 𝑄, 𝜋cf shortcut status (𝑄=+1: aligned; 𝑄=−1: conflicting) and the sample conflict fraction in Proposition 1

𝑏𝜉 , 𝑝𝜔 𝜙ref 𝜏 𝜇𝑐 Π(𝜙) 𝜆𝑡 tr , wg

observed training environments used by invariant learning ms set of milestone training steps 𝑇warm , 𝑇IRM lengths of the warm-up and invariant-training stages 𝑇proto , 𝑇eval intervals between prototype refreshes and validation evaluations 𝜃𝑡 , ℎ(𝑡) encoder retained at step 𝑡 and its repaired linear head

tr

Definition 1 (Endogenous environments with repair-aware selection (ERAS)). Given a checkpoint 𝜙̃ on the encoder trajectory, a pair (𝜙⋆ , 𝑡⋆ ) solves the ERAS problem when the induced partition, optimized representation, and retained candidates arise from the same trajectory and satisfy ̃ ̂ = Π(𝜙),

( )} 𝜙 ∈ arg min tr (𝜙) + 𝜆inv 𝜙, ̂ . ⋆

{

(5)

𝜙

and ranks the retained candidates after classifier repair, ( ) (6) 𝑡⋆ ∈ arg max WGAval ̂ ℎ𝑡 ◦𝜃𝑡 . 𝑡∶𝜃𝑡 ∈

The checkpoint 𝜙̃ lies on the same representation-learning trajectory subsequently optimized under its induced partition. Equation (6) ranks candidate encoders after fitting a matched deployment head to each one. Definition 1 aligns candidate generation with the learning trajectory and candidate selection with the repaired predictor used at deployment; neither condition requires training-group labels. The second condition is decisive because standard twostage pipelines rank candidates before repair, using the Stage 1 classifier 𝑓𝑡 : ( ) 𝑡pre ∈ arg max WGAval 𝑓𝑡 ◦𝜃𝑡 . (7) 𝑡∶𝜃𝑡 ∈

This criterion mixes representation quality with the bias of 𝑓𝑡 , and it can discard an encoder that would have repaired well. This motivates comparing candidates under the repaired classifier used at deployment. Q. Wang et al.: Preprint submitted to Elsevier

Explanation

The ideal endogenous objective would update Π(𝜙) with 𝜙. Equation (5) uses a trajectory checkpoint 𝜙̃ to make this coupling tractable during training.

4. Method In this section, we present ProME, a two-stage framework for group robustness without training-group labels. ProME instantiates the ERAS problem of Definition 1. It derives environments from the representation being trained, and it ranks candidates after classifier repair. Figure 2 separates candidate generation from deploymentaligned selection. Stage 1 performs a prototype-guided ERM warm-up, partitions the training data by the median prototype margin, and learns across the induced low- and high-margin environments. Checkpoints retained along this trajectory form the candidate set. Stage 2 removes each Stage 1 head, fits a matched group-balanced linear head to every frozen encoder, and ranks the repaired encoder–head pairs by validation WGA. The search is performed offline; test-time inference uses only the selected pair.

4.1. Cosine-Prototype Classification ProME requires an environment signal that follows the same representation geometry used for prediction. We adopt a cosine-prototype classifier to provide this shared geometric basis. Let 𝑏𝜉 denote the backbone and 𝑝𝜔 denote the projection layer. The normalized representation is defined as ( ) 𝜙𝜉,𝜔 (𝑥) = normalize 𝑝𝜔 (𝑏𝜉 (𝑥)) ,

(8)

Page 4 of 17

ProME: Group Robustness without Training Group Labels

Fig. 2. Overview of ProME. Stage 1 learns representation candidates from prototype-margin environments without traininggroup labels. Stage 2 repairs and selects candidates using group-annotated validation data, with five-fold cross-validation for regularization tuning. The dashed path denotes ProME-Refresh.

where normalize(𝑢) = 𝑢∕‖𝑢‖2 . For each class 𝑐, ProME computes the class prototype from the current representation: ) ( ∑ 1 𝜙(𝑥𝑖 ) , 𝜇𝑐 ← normalize |𝑐 | (𝑥 ,𝑦 )∈

means and terminal classifier geometry is further supported by neural-collapse results [28, 38]. The updated prototypes provide a geometric measure of label support for each example. This measure defines the partitioning signal used in the next module.

(9)

Without training-group labels, ProME constructs environments from the geometric signals available in the learned representation. Under the shortcut-first hypothesis, examples violating the dominant spurious correlation receive weaker support for their observed labels after ERM warmup. We quantify this support through the prototype margin:

𝑖

𝑖

𝑐

𝑐 = {(𝑥, 𝑦) ∈ tr ∶ 𝑦 = 𝑐}.

Given the two prototypes 𝜇0 and 𝜇1 , the prediction score is defined as ] 1[ 𝑓𝜙 (𝑥; 𝜇0 , 𝜇1 ) = cos(𝜙(𝑥), 𝜇1 ) − cos(𝜙(𝑥), 𝜇0 ) , (10) 𝜏 where 𝜏 > 0 is the temperature. The sign of 𝑓𝜙 (𝑥) determines the binary prediction, and its magnitude measures the difference in cosine support for the two classes. This makes the prediction rule directly interpretable through the feature– prototype geometry. Prototype updates. ProME treats prototypes as non-trainable state variables that are periodically refreshed during training. Every 𝑇proto steps, we update the prototypes using the closed-form class-mean rule in Eq. (9) and stop gradients through the prototype computation. During backpropagation, the refreshed prototypes remain fixed, and gradients update only 𝑏𝜉 and 𝑝𝜔 through 𝑓𝜙 . This construction follows the class-mean geometry of prototypical classifiers [37]. The alignment between normalized class Q. Wang et al.: Preprint submitted to Elsevier

4.2. Prototype-Margin Environment Inference

𝑠(𝑥; 𝑦, 𝜙, 𝜇0 , 𝜇1 ) = cos(𝜙(𝑥), 𝜇𝑦 ) − cos(𝜙(𝑥), 𝜇1−𝑦 ). (11) The margin measures the relative prototype support of the observed class. A larger value indicates stronger alignment with the observed label, while a smaller value indicates ambiguous or conflicting class support. The margin is computed from the same classifier used for prediction, which provides a unified geometric basis for environment inference. Environment construction. After prototype-guided ERM warm-up, we compute the margins over tr and construct two environments through a global median split: ̂0 = {(𝑥, 𝑦) ∶ 𝑠(𝑥; 𝑦) ≤ 𝑚}, Page 5 of 17

ProME: Group Robustness without Training Group Labels

̂1 = {(𝑥, 𝑦) ∶ 𝑠(𝑥; 𝑦) > 𝑚},

4.3. Classifier Repair and Candidate Selection (12)

𝑚 = mediantr 𝑠.

The median split produces two balanced cells up to ties without requiring group proportions. The low-margin environment contains examples with weaker prototype support for their observed labels, while the high-margin environment contains examples with stronger class agreement. Section 6.2.2 measures shortcut-conflict enrichment and downstream WGA. Corollary 1 separately uses the totalvariation quantity 𝜌 to state its transfer condition. ̂ Over the inferred enInvariant objective. Let 𝐾 = ||. vironments, we optimize the average environment risk together with the IRMv1 and REx penalties: ̄ 𝜙 ) = 𝐾 −1 (𝑓

∑

𝑒̂ (𝑓𝜙 ),

(13)

𝑒∈ ̂ ̂

The Stage 1 objective is ( ) ̄ 𝜙 ) + 𝜆𝑡 IRM (𝑓𝜙 ) + REx (𝑓𝜙 ) , (14) ProME (𝜙) = (𝑓 where 𝑤′ denotes the scalar dummy classifier in IRMv1, and 𝜆𝑡 controls the invariance regularization strength. The IRM penalty enforces consistent classifier optimality across environments, while the REx penalty reduces risk variation across inferred environments [45]. Section 5 bounds the gap between average and worst inferred-environment risk through the REx penalty (Proposition 3). IRMv1 adds a complementary constraint on environment-wise classifier optimality. Partition updates. Both ProME variants refresh the prototypes every 𝑇proto steps. The standard protocol constructs ̂ from the post-warm-up encoder and keeps the partition fixed during invariant training. ProME-Refresh updates ̂ by recomputing Eq. (12) after each prototype refresh. Both variants optimize the same objective and differ only in the update frequency of the representation-induced partition. Prototype refinement. For the label-augmented sensitivity analysis in Section 6.3, we replace Eq. (9) with groupbalanced prototypes: ) ( 1 1 𝜇𝑐bal ← normalize 𝜙̄ 𝑐,0 + 𝜙̄ 𝑐,1 , 2 2 ∑ 1 ̄ 𝜙(𝑥), (15) 𝜙𝑐,𝑎 = |𝑐,𝑎 | (𝑥,𝑦)∈ 𝑐,𝑎

where 𝑐,𝑎 = {(𝑥, 𝑦) ∈ tr ∶ 𝑦 = 𝑐, 𝐴(𝑥) = 𝑎}. This construction balances the contribution of spurious attributes within each class for sensitivity analysis. ProME uses class labels only to compute prototypes in the standard setting, following Eq. (9). Q. Wang et al.: Preprint submitted to Elsevier

(16)

 = {𝜃𝑡 ∶ 𝑡 ∈ ms } ∪ {𝜃best }.

𝑒∈ ̂ ̂

∑‖ | ‖2 IRM (𝑓𝜙 ) = 𝐾 −1 ‖∇𝑤′ 𝑒̂ (𝑤′ 𝑓𝜙 )| ′ ‖ , ‖ |𝑤 =1 ‖2 𝑒∈ ̂ ̂ ∑( ) ̄ 𝜙) 2 . 𝑒̂ (𝑓𝜙 ) − (𝑓 REx (𝑓𝜙 ) = 𝐾 −1

Stage 2 makes classifier repair the basis of representation selection. Given the retained backbones and groupannotated validation data, ProME fits the same groupbalanced linear-head family to every candidate and selects the repaired encoder–head pair with the highest validation WGA. Validation groups serve for candidate tracking, head fitting, hyperparameter tuning, and final model selection, without updating the Stage 1 representation. Candidate construction. Pre-repair performance is treated as one source of candidates, not as the deployment decision. ProME retains frozen backbones at fixed milestones ms and includes 𝜃best , the checkpoint with the highest prerepair validation WGA under the prototype classifier. All candidates are subsequently compared only after matched repair:

Each 𝜃𝑡 is a frozen backbone checkpoint at step 𝑡. Before classifier repair, ProME removes the Stage 1 projection and cosine-prototype head. This construction preserves representation diversity and evaluates all candidates under the same downstream classifier family. Candidate-specific classifier repair. For each 𝜃𝑡 ∈ , we freeze the encoder and fit a linear DFR head ℎ(𝑡) on the group-annotated validation set [15]. Let 𝑔val denote the validation examples in oracle group 𝑔 ∈ . For a regularization strength 𝛾, the group-balanced objective is (𝑡) (ℎ; 𝛾) = DFR

1 ∑ 1 || 𝑔∈ |𝑔 |

∑

val (𝑥,𝑦)∈𝑔 val

( ) 𝓁 ℎ(𝜃𝑡 (𝑥)), 𝑦 +𝛾‖ℎ‖22 . (17)

The objective assigns equal weight to each validation group regardless of its empirical frequency. For each candidate, we select 𝛾 by five-fold cross-validation and refit the head with the selected value. Freezing the encoder and using the same classifier family make the repaired candidates directly comparable. Validation selection. After fitting all candidate-specific heads, ProME evaluates each encoder–head pair using validation WGA: ( ) 𝑣𝑡 = WGAval ℎ(𝑡) ◦𝜃𝑡 ,

𝑡⋆ = arg max 𝑣𝑡 .

(18)

𝜃𝑡 ∈

⋆

The selected predictor is (𝜃𝑡⋆ , ℎ(𝑡 ) ). The group-annotated validation set is used for Stage 2 candidate tracking, head fitting, regularization tuning, and encoder–head ranking, whereas Stage 1 requires no training-group labels. Section 6.2.3 examines how candidate coverage and matched repair affect post-repair evaluation. Test-time prediction. Given a test input 𝑥, ProME applies the selected encoder and linear head: ) ⋆ ( 𝑦̂ = ℎ(𝑡 ) 𝜃𝑡⋆ (𝑥) . (19) Page 6 of 17

ProME: Group Robustness without Training Group Labels

Algorithm 1 ProME training and candidate selection. Input: tr , group-annotated val , 𝜏, {𝜆𝑡 }, 𝑇warm , 𝑇IRM , 𝑇proto , 𝑇eval , ms , and environment-update mode 𝗆𝗈𝖽𝖾. ⋆ Output: Selected encoder–head pair (𝜃𝑡⋆ , ℎ(𝑡 ) ). Stage 1: Prototype-Margin Environment Learning 1: Initialize backbone 𝑏𝜉 and projection 𝑝𝜔 . 2: Optimize (𝜉, 𝜔) for 𝑇warm epochs with 𝑓𝜙 . 3: {𝜇𝑐 } ← Ref resh(𝜙𝜉,𝜔 , tr ); ̂ ← Π(𝜙𝜉,𝜔 ). 4: (𝑢best , 𝜃best ) ← (WGAval (𝑓𝜙 ), copy(𝑏𝜉 )). 5: for 𝑡 = 1, … , 𝑇IRM do 6: if 𝑡 > 1 and 𝑡 ≡ 1 (mod 𝑇proto ) then 7: {𝜇𝑐 } ← Refresh(𝜙𝜉,𝜔 , tr ). 8: ̂ ← Π(𝜙𝜉,𝜔 ) if 𝗆𝗈𝖽𝖾 = Refresh. 9: end if ̂ 10: Update (𝜉, 𝜔) with Eq. (14) over . 11: if 𝑡 ≡ 0 (mod 𝑇eval ) then 12: 𝑢𝑡 ← WGAval (𝑓𝜙 ). 13: (𝑢best , 𝜃best ) ← (𝑢𝑡 , copy(𝑏𝜉 )) if 𝑢𝑡 > 𝑢best . 14: end if 15: if 𝑡 ∈ ms then 16: 𝜃𝑡 ← copy(𝑏𝜉 ). 17: end if 18: end for Stage 2: Classifier Repair and Candidate Selection 19:  ← {𝜃𝑡 ∶ 𝑡 ∈ ms } ∪ {𝜃best }. 20: for 𝜃𝑡 ∈  do 21: ℎ(𝑡) ← DFR5−fold (𝜃𝑡 , val ). 22: 𝑣𝑡 ← WGAval (ℎ(𝑡) ◦𝜃𝑡 ). 23: end for 24: 𝑡⋆ ← arg max𝜃𝑡 ∈ 𝑣𝑡 . ⋆ 25: return (𝜃𝑡⋆ , ℎ(𝑡 ) ).

We tune the regularization strength by five-fold crossvalidation on the group-annotated validation set, refit each candidate-specific head on the full set, and rank the repaired encoder–head pairs by validation WGA. Further implementation details are provided in Appendix B. Notably, ProME offers several practical advantages:

examples. Proposition 2 bounds how far the induced partition moves when the representation changes. Proposition 3 connects the REx term to the worst inferred-environment risk. The deployment-aligned selection condition is evaluated experimentally in Section 6. IRM and REx assume that the environments are already given. Identifying invariant structure from the pooled distribution of (𝑋, 𝑌 ) needs an additional inductive bias. ProME supplies that bias through prototype-margin geometry. Section 2 places the alternatives in context.

5.1. When prototype margins expose shortcut conflicts For binary labels, write 𝑌̃ = 2𝑌 − 1 ∈ {−1, +1}. Let 𝐴̃ ∈ {−1, +1} encode a binary spurious attribute and define ̃ Here 𝑄 = +1 denotes shortcut-aligned examples 𝑄 = 𝑌̃ 𝐴. and 𝑄 = −1 denotes shortcut-conflicting examples. This notation is used only for analysis. ProME never observes 𝐴 during Stage 1. The margin and the training logit must first be shown to share one direction. Lemma 1 (Prototype-margin identity). Under Assumption 1, let 𝑤⋆ = 𝜇1 − 𝜇0 . For every (𝑥, 𝑦), 𝑠(𝑥; 𝑦) = 𝑦̃ ⟨𝜙(𝑥), 𝑤⋆ ⟩,

𝑦̃ = 2𝑦 − 1.

(20)

Moreover, 𝑓𝜙 (𝑥) = 𝜏 −1 ⟨𝜙(𝑥), 𝑤⋆ ⟩. The environment score and the prediction logit share one linear direction. The proof is given in Appendix A.1. The identity alone does not imply shortcut-conflict separation. Proposition 1 gives sufficient conditions under which shortcut-conflicting examples receive lower margins. Proposition 1 (Shortcut-conflict separation). Suppose the signed representation after ERM warm-up satisfies 𝑌̃ 𝜙(𝑋) = 𝑢 + 𝑄𝑣 + 𝜉,

𝑢⊤ 𝑣 = 0.

(21)

Here, 𝑄 = +1 denotes shortcut alignment and 𝑄 = −1 denotes shortcut conflict. Let the prototype contrast be 𝑤⋆ = • Training-group-free candidate generation. Prototype𝛼𝑢 + 𝛽𝑣, where 𝛼, 𝛽 > 0. Assume ℙ(𝑄 = 𝑞) > 0 for each margin environments use class labels only and require 𝑞 ∈ {−1, +1}. If 𝔼[⟨𝜉, 𝑤⋆ ⟩ ∣ 𝑄] = 0, then no auxiliary environment model.

• Repair-aware representation selection. Every candidate is evaluated under the matched classifier family used for deployment. • Single-model deployment. Candidate search is offline; inference uses one selected encoder and one linear head. We summarize the complete ProME paradigm in Algorithm 1.

5. Theoretical Analysis This section analyzes the Stage 1 mechanism required by the first condition of Definition 1. Lemma 1 links the prototype margin to the classifier used for training. Proposition 1 states when low margins mark shortcut-conflicting Q. Wang et al.: Preprint submitted to Elsevier

𝔼[𝑠(𝑋; 𝑌 ) ∣ 𝑄 = +1] − 𝔼[𝑠(𝑋; 𝑌 ) ∣ 𝑄 = −1] = 2𝛽‖𝑣‖22 . (22) If, in addition, |⟨𝜉, 𝑤⋆ ⟩| ≤ 𝛿 almost surely for some 0 ≤ 𝛿 < 𝛽‖𝑣‖22 , every shortcut-conflicting example has a smaller margin than every shortcut-aligned example almost surely. For a sample of size 𝑛 with distinct margins, let 𝑛cf be its number of conflicting examples and 𝜋cf = 𝑛cf ∕𝑛. If 𝜋cf ≤ 1∕2, the low-margin cell contains all 𝑛cf conflicting examples and has conflict fraction 𝑛cf ∕⌈𝑛∕2⌉. This fraction equals 2𝜋cf when 𝑛 is even. The proof is given in Appendix A.2. The margin is informative when ERM learns a shortcut direction and the projected noise preserves its separation. Proposition 1 states as a condition what error- and bias-based methods observe empirically. Section 6.2.2 tests it on real data. Page 7 of 17

ProME: Group Robustness without Training Group Labels

5.2. Stability of the endogenous partition ProME computes both prototypes and the median on the training sample. A useful stability result must account for ̂ 𝑛 (𝜙) denote the empirical split this data dependence. Let Π in Eq. (12). Consistent with that equation, define 𝑒̂𝑖 (𝜙) = 0 when 𝑠𝑖 (𝜙) ≤ 𝑚̂ 𝑛 (𝜙) and 𝑒̂𝑖 (𝜙) = 1 otherwise. Proposition 2 (Same-sample partition stability). Fix a labeled sample tr = {(𝑥𝑖 , 𝑦𝑖 )}𝑛𝑖=1 with |𝑐 | > 0 for both classes. Let 𝜙 and 𝜙′ be unit-normalized representations whose class prototypes are computed from this same sample by Eq. (9). Define

The proof is given in Appendix A.4. Equation (25) is a deterministic inequality for a fixed predictor, partition, and vector of cell risks. Both terms on its right are aligned with the Stage 1 objective, but the inequality does not by itself bound the empirical-to-population gap or guarantee that optimization reaches a small right-hand side. It requires no calibration assumption linking the IRMv1 gradient penalty to worst-cell risk. IRMv1 remains a complementary optimization signal, while REx supplies the direct riskdispersion control. We next relate inferred-cell risk to oracle-group risk under an alignment condition.

𝜀 = max ‖𝜙(𝑥𝑖 ) − 𝜙′ (𝑥𝑖 )‖2 ,

Corollary 1 (Conditional transfer to oracle groups). Assume the loss takes values in [0, 𝐵]. Condition on the learned pre4 dictor and partition, and define TV(𝑃 , 𝑄) = sup |𝑃 (𝐴) − 𝐿𝑠 = 2 + . 𝐴 𝛽𝑛 𝑄(𝐴)| for distributions over (𝑋, 𝑌 ). If every oracle group distribution 𝑄𝑔 is within total variation 𝜌 of some inferredAssume 𝛽𝑛 > 0. Let 𝑚̂ 𝑛 (𝜙) be the empirical median and cell distribution 𝑄𝑒(𝑔) , then define the empirical concentration function √ () ̄ ) + (𝐾 − 1)REx (𝑓 ) + 𝐵𝜌. 𝑛 (26) (𝑓 ) ≤ (𝑓  ∑ [ ] wg 1 𝜔̂ 𝜙 (𝑡) = 1 |𝑠𝑖 (𝜙) − 𝑚̂ 𝑛 (𝜙)| ≤ 𝑡 . 𝑛 𝑖=1 The same appendix section proves the corollary. Corollary 1 controls oracle-group risk only under the stated alignThen the prototype, margin, and median perturbations satment condition. For the 0–1 loss, where √ 𝐵 = 1, it equivaisfy ̄ ) − (𝐾 − 1)REx (𝑓 ) − 𝜌. lently yields WGA(𝑓 ) ≥ 1 − (𝑓 2𝜀 The experiments measure shortcut-conflict enrichment; they , max ‖𝜇𝑐 (𝜙) − 𝜇𝑐 (𝜙′ )‖2 ≤ 𝑐 do not estimate 𝜌. 𝛽𝑛 Together, the results characterize three Stage 1 propermax |𝑠𝑖 (𝜙) − 𝑠𝑖 (𝜙′ )| ≤ 𝐿𝑠 𝜀, 𝑖 ties under explicit conditions: shortcut-conflict separation, |𝑚̂ 𝑛 (𝜙) − 𝑚̂ 𝑛 (𝜙′ )| ≤ 𝐿𝑠 𝜀. (23) same-sample partition stability, and fixed-risk control across inferred environments. Classifier repair and repair-aware Consequently, the fraction of reassigned examples obeys selection remain outside the analysis. 𝑛 ∑ 1 1[𝑒̂𝑖 (𝜙) ≠ 𝑒̂𝑖 (𝜙′ )] ≤ 𝜔̂ 𝜙 (2𝐿𝑠 𝜀). (24) 𝑛 𝑖=1 6. Experiments 𝑖

‖ ‖ ‖ ‖ 1 ∑ ‖ 𝜓(𝑥𝑖 )‖ 𝛽𝑛 = min ′ ‖ ‖ , 𝑐,𝜓∈{𝜙,𝜙 } ‖ |𝑐 | ‖ 𝑖∶𝑦𝑖 =𝑐 ‖2 ‖

The proof is given in Appendix A.3. Proposition 2 is deterministic and covers the implementation that ProME actually runs, in which prototypes and the median come from the same training set. A concentration argument with independently fixed prototypes would not. Equation (24) also identifies the failure modes. Instability increases when class means approach zero, which enlarges 𝐿𝑠 , or when many margins accumulate near the median, which enlarges 𝜔̂ 𝜙 .

5.3. Risk control on inferred environments Margin separation and partition stability explain how the environments are formed, but neither result alone guarantees robustness. We next characterize the quantity controlled directly by the REx penalty in Eq. (13). For a fixed classifier ̂ ) ̄ = 𝐾 −1 ∑ 𝑟𝑒 , and (wg 𝑓 , write 𝑟𝑒 = 𝑒̂ (𝑓 ),  = max𝑒 𝑟𝑒 . 𝑒

Proposition 3 (Worst inferred-environment risk bound). For any 𝐾 ≥ 2 inferred environments, √ ̂ ) ̄ ) + (𝐾 − 1)REx (𝑓 ). (25) (wg (𝑓 ) ≤ (𝑓 For the two-cell partition used by ProME, equality holds. Q. Wang et al.: Preprint submitted to Elsevier

In this section, we present the experimental results to evaluate two claims: (i) ProME improves worst-group accuracy without training-group labels; and (ii) prototype margins enrich shortcut-conflicting examples, while classifier repair reshapes candidate evaluation.

6.1. Setup Datasets. We evaluate ProME on four subpopulation-shift benchmarks, each pairing a label with a spurious attribute: Waterbirds [1], bird type against background; CelebA [5], blond hair against gender; CivilComments [46], toxicity across identity-based subpopulations; and ColoredMNIST [19], digit parity against a controllable color shortcut. The first three datasets form the main evaluation, while ColoredMNIST provides a controlled test under shortcut reversal. Baselines. We group the baselines by their access to group labels. When training-group labels are available, Group DRO [1] minimizes the worst-group loss, while RWG [20] balances group sampling. Without validation-group access, ERM [47] minimizes average loss, while XRM+GroupDRO [7] infers groups without oracle group labels. EIIL [9], JTT [6], CnC [11], and AFR [12] use validation groups only for Page 8 of 17

ProME: Group Robustness without Training Group Labels

model selection. When validation groups can also be used for fitting, SSA [48] infers training-group labels, MAPLE [49] learns weights for training examples, and DFR [15] fits a balanced classifier on a frozen representation. GSR-HF and GSR [50] use influence scores to reweight examples for lastlayer retraining. This grouping makes the supervision budget of each comparison explicit.

WGA, compared with 86.76 ± 2.30% for the ERM-loss control (Table 3). This auxiliary comparison complements the class-only conflict-enrichment evidence in Fig. 3. The result is consistent with the conditional separation in Proposition 1. Prototype margins also reach 75.77% on ColoredMNIST, compared with 74.09% for ERM loss. In the label-augmented CelebA diagnostic, the two pipelines yield Implementation and Evaluation. We use ImageNet-pretrained means that differ by 0.56 points, with a shared standard deviation of 0.94 points. Section 6.2.3 examines candidate ResNet-50 encoders for Waterbirds and CelebA, BERTcoverage on CelebA. base-uncased for CivilComments, and the three-layer conTrajectory-derived margins improve pre-repair WGA volutional network of IRM [19] for ColoredMNIST. The in the matched source comparison. The live-margin and candidate pool contains fixed milestone checkpoints and frozen-reference variants use the same margin criterion and the validation-best pre-repair encoder. We freeze each caninvariance weight, but derive margins from different endidate and fit a group-balanced logistic-regression head, coders (Table 4). The live source raises pre-repair WGA selecting its 𝓁2 penalty by five-fold cross-validation. We from 68.95 ± 3.39% to 78.85 ± 3.56%. After matched clasreport WGA over the oracle test groups. Waterbirds, CelebA, sifier repair, their final WGA values are 90.34 ± 1.75% and and ColoredMNIST contain four (𝑦, 𝑎) groups, while Civil90.60 ± 0.95%, respectively. This matched comparison tests Comments uses the standard 16 identity–label groups. The the effect of the partition source before repair. Proposition 2 headline results report the mean and population standard separately bounds assignment changes as the representation deviation over seeds {42, 0, 1}, and we also evaluate a fixed evolves. pool of ten seeds. Table 2 uses published baseline results The invariance penalty affects pre-repair WGA, while from their original protocols, while the source, penalty, the global median matches the class-wise diagnostic. Reand candidate-set comparisons use matched settings; labelmoving the penalty lowers pre-repair WGA from 78.85 ± augmented diagnostics are identified explicitly. Appendix B 3.56% to 72.95 ± 1.02% (Table 4). Their final WGA values provides the remaining training details. Because baseline after matched repair are 90.34 ± 1.75% and 90.71 ± 1.66%, entries follow their published protocols, Table 2 provides respectively. The observed penalty effect is concentrated reported rather than controlled comparisons. in pre-repair WGA. Replacing the global median with a class-wise median produces similar environment composi6.2. Main Results tion (0.68 and 0.02) and conflict balanced accuracy (0.74 ± 6.2.1. Overall Group Robustness 0.01 and 0.74 ± 0.02). The global rule provides comparable ProME achieves the highest average WGA among methdiagnostic separation with one threshold. ods compared at the same group-label access. ProME uses validation groups for classifier fitting and selection, 6.2.3. Classifier Repair Reshapes Candidate matching the access of SSA, MAPLE, DFR, GSR-HF, and Evaluation GSR. Its average WGA is 87.0%, 3.1 points above GSR The second ERAS condition ranks encoder–head pairs (83.9%), the highest baseline entry at this access level. by validation WGA after matched classifier repair (DefiniProME ranks first on Waterbirds (93.1%) and CivilComtion 1). We examine how candidate coverage and classifier ments (78.7%), and second on CelebA (89.3±0.5%). The repair affect post-repair evaluation. CivilComments result is 6.5 points above the strongest prior Multi-candidate selection improves post-repair WGA on entry in the table (72.2%). These differences compare pubCelebA. Fig. 5(c) holds the repair procedure fixed and varies lished results under their original protocols. The following only the candidate set. A single retained checkpoint reaches analyses use matched settings to examine the two alignment 86.94% WGA after repair. Ranking repaired candidates from mechanisms. milestone and random pools reaches 89.31% and 89.87%,

6.2.2. Prototype Margins Enrich Shortcut Conflicts Definition 1 draws the partition from a checkpoint on the trajectory it later regularizes. This subsection examines the induced environments through conflict enrichment and downstream controls over the partition signal, its source encoder, and the invariance penalty. Prototype margins concentrate shortcut conflicts in the low-margin environment. Under the global-median rule, shortcut-conflicting examples constitute 0.68 of the lowmargin environment and 0.01 of the high-margin environment (Fig. 3). Fig. 4 shows the same ordering at the example level. In the label-augmented Waterbirds diagnostic, the prototype-margin pipeline reaches 90.31 ± 1.29% Q. Wang et al.: Preprint submitted to Elsevier

respectively. Both pools outperform the single-checkpoint setting in these runs. Classifier repair narrows performance differences among Stage 1 variants. Across four Waterbirds variants, prerepair WGA spans 12.20 points, whereas matched repair reduces the final spread to 1.25 points (Table 4). CivilComments shows the same pattern under the shortened schedule in Fig. 5(b), where the spread contracts from 9.98 to 0.34 points. These contractions show that comparisons based on Stage 1 heads need not persist after repair. Repair-aware selection instead evaluates the encoder–head pairs used at deployment.

Page 9 of 17

ProME: Group Robustness without Training Group Labels Table 2 Performance comparison on three real-world benchmarks. For ProME, we report the mean and standard deviation (± std) of WGA over seeds {42, 0, 1}; baseline entries are taken from their original protocols. Train indicates whether oracle training-group labels update the representation, and Val indicates oracle validation-group use: S for selection or hyperparameter tuning, F+S for classifier fitting and selection, and – for none. Rows are grouped by increasing group-label access. Average is the unweighted mean across datasets. bolded and underlined values mark the highest and second-highest reported means. Group-label

Method

Train

Worst-group accuracy (%)

Val

Waterbirds

CelebA

CivilComments

Average

Training-group supervision (reference) Group DRO [1] Yes RWG [20] Yes

S S

91.4±1.1 87.6±1.6

88.9±2.3 84.3±1.8

70.0±2.0 72.0±1.9

83.4 81.3

No group labels ERM [47] XRM+GroupDRO [7]

No No

– –

74.9±2.4 88.1

46.9±2.8 89.1

55.6±0.6 72.2

59.1 83.1

Validation groups for selection EIIL [9] No JTT [6] No CnC [11] No AFR [12] No

S S S S

78.7 86.7 88.5±0.3 90.4±1.1

83.3 81.1 88.8±0.9 82.0±0.5

67.0 69.3 68.9±2.1 68.7±0.6

76.3 79.0 82.1 80.4

Validation groups for fitting and selection SSA [48] No MAPLE [49] No DFR [15] No GSR-HF [50] No GSR [50] No ProME (ours) No

F+S F+S F+S F+S F+S F+S

89.0±0.6 91.7 92.9±0.2 87.5±0.1 92.9±0.0 𝟗𝟑.𝟏±0.3

𝟖𝟗.𝟖±1.3 88.0 88.3±1.1 86.3±0.4 87.0±0.4 89.3±0.5

69.9±2.0 64.1 70.1±0.8 68.9±0.2 71.7±0.6 𝟕𝟖.𝟕±0.2

82.9 81.3 83.8 80.9 83.9 𝟖𝟕.𝟎

Oracle groups

(a)

Global median (ProME)

(b)

Class-wise diagnostic

Ten-seed diagnostic

(c)

0.80

(d)

Global: 0.74 ± 0.02 Class-wise: 0.74 ± 0.01

Conflict BA

0.75

Group fractions 0.4

0.65 0.60

0.38 0.41

0.55

0.2 0

0.70

0.11 0.11

LL

LW WL WW

P(conflict ∣ low) = 0.68 P(conflict ∣ high) = 0.01

P(conflict ∣ low) = 0.68 P(conflict ∣ high) = 0.02

0.50 Global median

Landbird/land (LL)

Landbird/water (LW)

Waterbird/land (WL)

Waterbird/water (WW)

high margin

low margin

Class-wise median

prototype

Fig. 3. Prototype-margin enrichment of shortcut conflicts on Waterbirds. Panels (a–c) show the same seed-42 test examples in a shared t-distributed stochastic neighbor embedding (t-SNE). (a) Oracle groups; the inset reports their fractions in the displayed subset. (b) Environments induced by the global-median rule used by ProME. (c) Environments induced by a class-wise median diagnostic. Marker shapes denote oracle groups, colors denote inferred environments, and stars denote class prototypes. (d) Conflict balanced accuracy on the training split over ten seeds. Gray lines connect the same seed; diamonds and error bars show the mean and standard deviation.

Classifier repair reduces a posterior-alignment diagnostic on ColoredMNIST. Fig. 6 compares empirical ℙ(𝑦=1 ∣ ℎ, 𝑒) across two training environments and the reversed test environment. ProME has the largest integrated gap before repair and the smallest gap after repair, below ERM and the ERM-loss control. he posterior-alignment gap characterizes Q. Wang et al.: Preprint submitted to Elsevier

how classifier repair changes score geometry under shortcut reversal. Corollary 1 states the transfer condition in terms of the total-variation quantity 𝜌.

Page 10 of 17

Lowest margin

ProME: Group Robustness without Training Group Labels

WL s = − 0.99

WL s = − 0.98

WL s = − 0.96

LW s = − 0.96

WL s = − 0.95

LW s = − 0.95

WL s = − 0.95

LL s = 1.02

LL s = 1.01

LL s = 1.01

LL s = 1.01

LL s = 1.01

LL s = 1.01

LL s = 1.01

LL s = 1.01

Highest margin

WL s = − 1.00

Fig. 4. Waterbirds examples ranked by prototype margin. The displayed lowest-margin examples contain label–background conflicts, whereas the displayed highest-margin examples show aligned pairs. Oracle group labels are included only to support interpretation of the examples. (b) CivilComments

90 85 80 source effect +9.90

75 70 65

Pre-repair

(c) CelebA candidate coverage

spread 9.98 → 0.34

DFR WGA (\%) on CelebA

spread 12.20 → 1.25

Worst-group accuracy (\%)

Worst-group accuracy (\%)

(a) Waterbirds 95

70

60

50

40

Final

Pre-repair

Live margin

Live loss

Final No penalty

91 89.87

90 89.31

89 88 86.94

87 86

Single-ckpt Milestone Random-K DFR (K = 4) (K = 5) Frozen reference

Fig. 5. Repair-aware candidate evaluation. (a) Waterbirds and (b) CivilComments before and after classifier repair; each line connects one Stage 1 variant of Table 4. (c) CelebA with single-checkpoint, milestone, and random candidate sets under a fixed repair procedure; the band marks the full multi-candidate pipeline.

Table 3 Downstream WGA for prototype-margin and ERM-loss partitions. This is a pipeline-level comparison: each partition passes through the same Stage 2 repair and selection protocol. For the prototype-margin row, Waterbirds and CelebA are labelaugmented diagnostics using the group-balanced prototypes of Eq. (15), reported over five and three seeds; ColoredMNIST follows the class-only protocol over ten seeds. Partition

Worst-group accuracy (%) Waterbirds

CelebA

ERM-loss 86.76±2.30 85.93±0.94 Prototype-margin 90.31±1.29 85.37±0.94

ColoredMNIST

Variant

Partition Source Criterion

𝜆

Worst-group accuracy (%) Pre-repair

Live margin Live Margin 3 78.85±3.56 Live loss Live Loss 3 81.15±1.11 No penalty Live Margin 0 72.95±1.02 Frozen reference Frozen Margin 3 68.95±3.39

Final 90.34±1.75 91.59±1.21 90.71±1.66 90.60±0.95

74.09±5.75 75.77±5.48

6.3. Seed Budget, Label Budget, and Hyperparameters Real-world WGA remains stable across the fixed seed budgets. Across ten fixed seeds, Waterbirds, CelebA, and CivilComments have standard deviations of 0.62, 0.51, and Q. Wang et al.: Preprint submitted to Elsevier

Table 4 Waterbirds WGA before and after classifier repair. Each variant changes one element of Stage 1 and then passes through the same Stage 2. Source is the encoder the partition is read from. Mean ± standard deviation over three seeds.

0.30 WGA points (Fig. 7). Their ten-seed means remain within 0.62 points of the three-seed results in Table 2. Fig. 9(d) shows little change as the seed budget grows from 3 to 10. ColoredMNIST has a wider standard deviation of 5.48 points, so the stability conclusion is limited to the three real-world benchmarks.

Page 11 of 17

(a) ERM

1.0

(b) ERM-loss

1.0

0.5

0.5

0.0 −2.0

−1.5

−1.0

−0.5 0.0 0.5 Classifier logit h

1.0

(c) ProME

1.0

0.5

0.0 −1.5

−1.0

−0.5 0.0 0.5 Classifier logit h

0.0 −1.5

1.5 Integrated posterior-alignment gap

Empirical ℙ(Y = 1 ∣ h, e)

Empirical ℙ(Y = 1 ∣ h, e)

ProME: Group Robustness without Training Group Labels

1.0

Train env. 1

−1.0

−0.5 0.0 0.5 Classifier logit h

1.0

(d) Alignment gap lower is better

0.6 0.4 0.2 0.0

Train env. 2

ERM

ERM-loss

ProME pre-repair

ProME final

Reversed test

Fig. 6. Classifier repair reduces the posterior-alignment gap on ColoredMNIST. Panels (a)–(c) plot the empirical ℙ(𝑦=1 ∣ ℎ, 𝑒) for two training environments and the reversed test environment under ERM, ERM-loss, and ProME after classifier repair. Panel (d) reports the integrated posterior-alignment gap for ERM, ERM-loss, and ProME before and after repair. A lower value indicates better alignment. Gray lines connect the same seed. Curves use 50 equal-width bins with at least five samples per bin.

92:48 §0:62

worst-group accuracy (\%)

95

median / IQR per-seed run mean

89:16 §0:51

90

75:77 §5:48

85 78:42 §0:30

80

Table 5 Effect of the data used to fit the classifier-repair head. The ProME representation is fixed; only the group-balanced fitting data changes. Mean ± standard deviation over ten seeds. 𝑁 is the training-set size and 𝑛min its smallest (𝑦, 𝑎) group. Smallest group 𝑛min

train-DFR (%)

val-DFR (%)

Waterbirds (4,795)

56

73.80±2.51

92.48±0.62

CelebA (162,770)

1,387

87.61±1.00

89.44±0.82

Dataset (𝑁)

75 70

bird

er Wat

s

ebA

Cel

ent

Civ

mm ilCo

s

T

NIS

dM lore

Co

Fig. 7. WGA across ten fixed seeds. Boxes show interquartile ranges, center lines show medians, points show individual runs, and diamonds show means.

The value of validation-group labels depends on minoritygroup support. Stage 2 uses group-annotated validation data for classifier repair. Table 5 holds the ProME representation fixed and changes only the group-balanced data used to fit the head. Val-DFR exceeds train-DFR by 18.68 points on Waterbirds, whose smallest training group contains 56 Q. Wang et al.: Preprint submitted to Elsevier

examples. The gap is 1.83 points on CelebA, whose smallest group contains 1,387 examples. Across six ColoredMNIST settings, the train-DFR minus val-DFR gap approaches zero as 𝑛min increases (Fig. 8). These results support validationgroup labels for classifier repair when training minority support is limited. The source and partition controls above examine the Stage 1 environment mechanism. Stage 1 representation learning uses no training-group labels. The label-augmented Waterbirds diagnostic is stable across the swept settings. Fig. 9(a–c) sweeps the invariance weight, the temperature, and the prototype-refresh period on Waterbirds under the label-augmented refinement with group-balanced prototypes. Every swept value of all three hyperparameters stays inside a 0.93-point band, from 92.14% to 93.07%, and this includes the endpoints 𝜆 = 30 Page 12 of 17

Δ = train-DFR − val-DFR (WGA %)

ProME: Group Robustness without Training Group Labels 2 → 0 (large groups)

0 -2 -4 -6

penalty (small groups)

-8 71

143

357

715

1787

3575

Smallest training group nmin

Fig. 8. Effect of minority support on classifier-repair label source. train-DFR minus val-DFR WGA versus the smallest training group 𝑛min , reported as mean ± standard deviation over five seeds.

WGA (%)

(a)

(b)

94

94

93

93

92

92

91

91 1

3

λ

10

30

WGA (%)

(c)

CRediT authorship contribution statement

0.050.07 0.1 0.150.2

τ

(d)

94

94

93

86

92

78

91 100

500

Tproto Waterbirds

CelebA

∞

Qianqian Wang: Conceptualization, Methodology, Software, Validation, Formal analysis, Investigation, Writing – original draft. Yunshan Li: Writing – review & editing. Dawei Huang: Software, Resources. Wenwu Gong: Writing – review & editing, Methodology, Conceptualization. Lili Yang: Supervision, Resources, Writing – review & editing.

Declaration of competing interest The authors declare that they have no known competing financial interests or personal relationships that could have appeared to influence the work reported in this paper.

70 50

formulating group robustness without training-group labels as the endogenous environments with repair-aware selection (ERAS) problem, ProME couples trajectory-based environment construction with post-repair comparison of encoder– head pairs. Theoretically, we characterize shortcut-conflict separation and partition stability under explicit conditions, and bound worst inferred-environment risk for a fixed predictor and partition, with transfer to oracle groups under total-variation alignment. Empirically, ProME achieves the highest average WGA among methods under matched group-label access, while analyses show pre-repair gains from trajectory-derived margins and reduced differences among Stage 1 variants after classifier repair. As a backboneagnostic framework that learns representations without traininggroup labels and deploys a single encoder–head pair, ProME offers deployment alignment as a practical design principle for group-robust learning, and we hope this work facilitates future research across diverse forms of subpopulation shift.

3

5

7

10

Number of seeds CivilComments

ColoredMNIST

Fig. 9. Hyperparameter and seed-budget sensitivity. (a– c) Label-augmented Waterbirds diagnostic across the invariance weight, temperature, and prototype-refresh period; bands show standard deviations and dashed lines mark defaults. (d) WGA across seed budgets.

Acknowledgements This research was funded by the SUSTech Presidential Postdoctoral Fellowship, the China Postdoctoral Science Foundation (Grant No. 2025M773057), Shenzhen Science and Technology Program (Grant No. ZDSYS20210623092007023), and the Shenzhen Key Laboratory of Safety and Security for Next Generation of Industrial Internet, Southern University of Science and Technology.

Data availability at ten times the default weight and 𝑇proto = ∞, which suspends prototype refreshing altogether. Within this Waterbirds sweep, WGA is insensitive to an order-of-magnitude change in the Stage 1 invariance weight.

7. Conclusion In this paper, we presented ProME, a two-stage and backbone-agnostic framework designed to align environment construction and model selection with the deployed predictor when training-group labels are unavailable. By Q. Wang et al.: Preprint submitted to Elsevier

The three real-world datasets are publicly available from their cited sources. ColoredMNIST is generated from MNIST using the protocol in Section 6. Our code will be publicly available upon acceptance.

Declaration of generative AI use The authors used generative AI tools only for language polishing and readability improvement. All technical ideas, experiments, analyses, and conclusions were developed and verified by the authors, who take full responsibility for the content of this paper. Page 13 of 17

ProME: Group Robustness without Training Group Labels

A.3. Proof of Proposition 2

A. Proofs for Theoretical Analysis This appendix provides the complete proofs for Section 5. The assumptions are local to the corresponding results and do not alter the algorithm in Section 4. Lemma 1 and Proposition 1 are conditional structural results; Proposition 2 is deterministic on a fixed sample; and Proposition 3 is an algebraic bound for fixed cell risks. Corollary 1 adds a conditional population-risk transfer under total-variation alignment. No finite-sample empirical-to-population guarantee is claimed. Assumption 1 (Prototype geometry). The representation is sphere-normalized, ‖𝜙(𝑥)‖2 = 1. The labeled sample used to form the prototypes contains both classes, and each empirical class mean is non-zero: 1 ∑ 𝜙(𝑥𝑖 ), 𝜇𝑐 = 𝜙̄ 𝑐 ∕‖𝜙̄ 𝑐 ‖2 . 𝜙̄ 𝑐 = |𝑐 | 𝑖∶𝑦 =𝑐 𝑖

The two empirical prototypes are distinct and the temperature satisfies 𝜏 > 0.

Proof. Write 𝜙̄ 𝑐 = |𝑐 |−1

𝑓𝜙 (𝑥) = 𝜏 −1 ⟨𝜙(𝑥), 𝜇1 − 𝜇0 ⟩ = 𝜏 −1 ⟨𝜙(𝑥), 𝑤⋆ ⟩. If 𝑦 = 1, the margin is ⟨𝜙(𝑥), 𝜇1 − 𝜇0 ⟩. If 𝑦 = 0, it is ⟨𝜙(𝑥), 𝜇0 − 𝜇1 ⟩. Since 𝑦̃ = +1 in the first case and −1 in the second, both cases are summarized by Eq. (20).

A.2. Proof of Proposition 1 Proof. Lemma 1 and Eq. (21) give 𝑠(𝑋; 𝑌 ) = ⟨𝑌̃ 𝜙(𝑋), 𝑤⋆ ⟩ = 𝛼‖𝑢‖22 + 𝑄𝛽‖𝑣‖22 + ⟨𝜉, 𝑤⋆ ⟩, where the last equality uses 𝑢⊤ 𝑣 = 0. Taking the conditional expectation and using 𝔼[⟨𝜉, 𝑤⋆ ⟩ ∣ 𝑄] = 0 yields 𝑞 ∈ {−1, +1}.

Subtracting the two conditional means proves Eq. (22). On the probability-one event where the bounded projectednoise condition holds, every conflicting example satisfies 𝑠 ≤ 𝛼‖𝑢‖22 − 𝛽‖𝑣‖22 + 𝛿, whereas every aligned example satisfies 𝑠 ≥ 𝛼‖𝑢‖22 + 𝛽‖𝑣‖22 − 𝛿. The second lower bound exceeds the first upper bound by 2(𝛽‖𝑣‖22 − 𝛿) > 0. Thus all conflicting margins precede all aligned margins in the empirical ordering almost surely. With distinct margins, the conventional empirical median assigns exactly ⌈𝑛∕2⌉ observations to the low-margin cell. If 𝑛cf ≤ 𝑛∕2, all conflicting observations lie in that cell, whose conflict fraction is 𝑛cf ∕⌈𝑛∕2⌉; for even 𝑛, this reduces to 2𝜋cf . Q. Wang et al.: Preprint submitted to Elsevier

𝜙̄ ′𝑐 = |𝑐 |−1

∑

𝜙′ (𝑥𝑖 ).

𝑖∶𝑦𝑖 =𝑐

The definition of 𝜀 and the triangle inequality give (A.1)

‖𝜙̄ 𝑐 − 𝜙̄ ′𝑐 ‖2 ≤ 𝜀. For non-zero vectors 𝑎 and 𝑏, ‖ 𝑎 2‖𝑎 − 𝑏‖2 𝑏 ‖ ‖ ‖ ‖ ‖𝑎‖ − ‖𝑏‖ ‖ ≤ min(‖𝑎‖ , ‖𝑏‖ ) . ‖ 2 2 ‖2 2 2

(A.2)

Applying Eq. (A.2) to the two empirical class means and using Eq. (A.1) proves (A.3)

max ‖𝜇𝑐 (𝜙) − 𝜇𝑐 (𝜙′ )‖2 ≤ 2𝜀∕𝛽𝑛 . 𝑐

For any example and either class 𝑐, adding and subtracting ⟨𝜙′ (𝑥𝑖 ), 𝜇𝑐 (𝜙)⟩ gives |⟨𝜙(𝑥𝑖 ), 𝜇𝑐 (𝜙)⟩ − ⟨𝜙′ (𝑥𝑖 ), 𝜇𝑐 (𝜙′ )⟩| ≤ ‖𝜙(𝑥𝑖 ) − 𝜙′ (𝑥𝑖 )‖2 + ‖𝜇𝑐 (𝜙) − 𝜇𝑐 (𝜙′ )‖2 ≤ (1 + 2∕𝛽𝑛 )𝜀. The margin is the difference of two such terms. Therefore max |𝑠𝑖 (𝜙) − 𝑠𝑖 (𝜙′ )| ≤ (2 + 4∕𝛽𝑛 )𝜀 = 𝐿𝑠 𝜀. 𝑖

(A.4)

Set 𝜂 = 𝐿𝑠 𝜀. A coordinatewise perturbation of a finite sequence by at most 𝜂 changes each order statistic by at most 𝜂. The same bound applies to the empirical median, whether defined as the central order statistic or the average of the two central order statistics. Equation (A.4) gives (A.5)

|𝑚̂ 𝑛 (𝜙) − 𝑚̂ 𝑛 (𝜙′ )| ≤ 𝜂.

= ⟨𝑢 + 𝑄𝑣 + 𝜉, 𝛼𝑢 + 𝛽𝑣⟩

𝔼[𝑠 ∣ 𝑄 = 𝑞] = 𝛼‖𝑢‖22 + 𝑞𝛽‖𝑣‖22 ,

𝜙(𝑥𝑖 ),

𝑖∶𝑦𝑖 =𝑐

A.1. Proof of Lemma 1 Proof. Because 𝜙(𝑥) and both prototypes have unit norm, every cosine in Eqs. (10) and (11) is an inner product. Hence

∑

Equations (A.3)–(A.5) prove Eq. (23). It remains to characterize which assignments can change. Consider an example assigned to the low-margin cell under 𝜙 and to the high-margin cell under 𝜙′ . Equations (A.4) and (A.5) imply 𝑚̂ 𝑛 (𝜙) − 2𝜂 < 𝑠𝑖 (𝜙) ≤ 𝑚̂ 𝑛 (𝜙). The reverse reassignment similarly implies 𝑚̂ 𝑛 (𝜙) < 𝑠𝑖 (𝜙) ≤ 𝑚̂ 𝑛 (𝜙) + 2𝜂. Thus every reassigned example lies in {|𝑠𝑖 (𝜙) − 𝑚̂ 𝑛 (𝜙)| ≤ 2𝜂}. Summing its indicator over the fixed sample and dividing by 𝑛 proves Eq. (24).

A.4. Proof of Proposition 3 and Corollary 1 ̄ choose one 𝑒⋆ ∈ Proof of Proposition 3. Let 𝑑𝑒 = 𝑟𝑒 − , ∑ arg max𝑒 𝑑𝑒 , and set 𝑀 = 𝑑𝑒⋆ . Since 𝑒 𝑑𝑒 = 0, the 𝐾 − 1 deviations indexed by 𝑒 ≠ 𝑒⋆ sum to −𝑀. By Cauchy– Schwarz, ( )2 ∑ ∑ 1 𝑀2 𝑑𝑒2 ≥ 𝑑𝑒 = . 𝐾 − 1 𝑒≠𝑒⋆ 𝐾 −1 𝑒≠𝑒⋆ Page 14 of 17

ProME: Group Robustness without Training Group Labels Table B.1 ColoredMNIST under two evaluation protocols. Published ERM and oracle-IRM results use reversed-environment test accuracy; the remaining rows use WGA under the ProME protocol. Results are mean ± standard deviation over ten seeds. Method

Seeds

Metric

Score (%)

ERM (published) IRM, oracle environments (published)

10 10

Test accuracy Test accuracy

17.1±0.6 66.9±2.5

JTT DFR Loss-split (EIIL-style) ProME

10 10 10 10

WGA WGA WGA WGA

68.14±2.09 70.86±0.55 74.09±5.75 75.77±5.48

Note. Published ERM and IRM values are quoted from Table 2 of [19]. The WGA block follows Section 6; its loss-split row changes the partition criterion within the same pipeline.

Consequently, REx =

1 ∑ 2 1 𝑑 ≥ 𝐾 𝑒 𝑒 𝐾

(

𝑀2 +

𝑀2 𝐾 −1

) =

𝑀2 . 𝐾 −1

√ ̄ proves (𝐾 − 1)REx , and adding  Therefore 𝑀 ≤ Eq. (25). For 𝐾 = 2, the two centered risks are 𝑑 and −𝑑, so REx = 𝑑 2 and the bound is an equality. Proof of Corollary 1. Under the convention TV(𝑃 , 𝑄) = sup𝐴 |𝑃 (𝐴) − 𝑄(𝐴)|, the variational characterization for a loss in [0, 𝐵] gives |𝔼𝑄𝑔 𝓁 − 𝔼𝑄𝑒(𝑔) 𝓁| ≤ 𝐵 TV(𝑄𝑔 , 𝑄𝑒(𝑔) ) ≤ 𝐵𝜌. Hence 𝑔 (𝑓 ) ≤ 𝑒(𝑔) (𝑓 ) + 𝐵𝜌. Maximizing over oracle groups gives ̂

() () wg (𝑓 ) ≤ wg (𝑓 ) + 𝐵𝜌.

Applying Proposition 3 to the first term proves Eq. (26).

B. Supplementary Experiments and Tables B.1. ColoredMNIST Protocol Comparison Table B.1 separates two evaluation conventions. Published ERM and oracle-IRM results use reversed-environment test accuracy. The matched block uses WGA under the ProME protocol and supports the diagnostic discussion in Section 6.2.2.

B.2. Single-Seed Pipeline Decomposition Table B.2 provides a seed-42 decomposition under a fixed budget and candidate set. Classifier repair and repairaware selection produce the largest change in all displayed arms. Replacing the partition also changes the final result, while removing the invariance penalty has a smaller effect in this run. Because the table contains one seed, it is used as a mechanism diagnostic rather than a population estimate. The multi-seed controls in Table 4 provide the main evidence.

Q. Wang et al.: Preprint submitted to Elsevier

Table B.2 Single-seed Waterbirds pipeline decomposition. All variants use class-only prototypes and the same training budget, candidate set, and validation protocol. Pre-repair evaluates Stage 1; Final includes classifier repair and selection; Gain = Final − Pre-repair. Variant

Partition

ProME (full) Prototype margin ERM-loss control ERM loss Without invariance Prototype margin

Invariance DFR Pre-repair Final ✓ ✓ ✗

✓ ✓ ✓

72.27 85.36 72.90

Gain

91.74 +19.47 86.29 +0.93 91.28 +18.38

B.3. Training Details Training schedules. The ERM warm-up lasts 5 epochs on the real-world benchmarks and 100 epochs on ColoredMNIST. Invariant training then runs for 3,000 steps on Waterbirds and CivilComments, 5,000 on CelebA, and 1,000 on ColoredMNIST. AdamW uses learning rates of 10−4 for ResNet-50 and 2 × 10−5 for BERT-base, with weight decay 10−4 and 10−2 , respectively. Batch sizes are 32, 64, and 16 on Waterbirds, CelebA, and CivilComments. The penalty weight increases linearly from 1 to 3 on the vision benchmarks and to 10 on CivilComments. Prototypes are refreshed every 𝑇proto = 100 steps by default, while the three-seed Waterbirds headline runs use 𝑇proto = 50. The training loss is logistic cross-entropy. Loss boundedness is required only for the optional total-variation transfer in Corollary 1. The controlled analyses in Sec. 6.2.2 and Sec. 6.2.3 run repeated seeds and repeated repair passes over each candidate pool, so we run the seed-heavy sweeps on Waterbirds and ColoredMNIST and reserve CelebA, whose 162,770 training images make each pass an order of magnitude more expensive, for the settings where it carries the comparison.

References [1] S. Sagawa*, P. W. Koh*, T. B. Hashimoto, P. Liang, Distributionally robust neural networks, in: International Conference on Learning Representations, 2020. URL https://openreview.net/forum?id=ryxGuJrFvS [2] Y. Yang, H. Zhang, D. Katabi, M. Ghassemi, Change is hard: A closer look at subpopulation shift, in: A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, J. Scarlett (Eds.), Proceedings of the 40th International Conference on Machine Learning, Vol. 202 of Proceedings of Machine Learning Research, PMLR, 2023, pp. 39584–39622. URL https://proceedings.mlr.press/v202/yang23s.html [3] W. Ye, L. Jiang, E. Xie, G. Zheng, Y. Ma, X. Cao, D. Guo, D. Qi, Z. He, Y. Tian, C. W. Porter, M. Coffee, Z. Zeng, S. Li, Z. Wang, T.-H. K. Huang, J. M. Rehg, H. Kautz, A. Zhang, The clever hans mirage: A comprehensive survey on spurious correlations in machine learning, Transactions on Machine Learning Research (2026). URL https://openreview.net/forum?id=kIuqPmS1b1 [4] R. Geirhos, J.-H. Jacobsen, C. Michaelis, R. Zemel, W. Brendel, M. Bethge, F. A. Wichmann, Shortcut learning in deep neural networks, Nature Machine Intelligence 2 (11) (2020) 665–673. doi: 10.1038/s42256-020-00257-z. URL https://doi.org/10.1038/s42256-020-00257-z [5] Z. Liu, P. Luo, X. Wang, X. Tang, Deep learning face attributes in the wild, in: 2015 IEEE International Conference on Computer Vision (ICCV), 2015, pp. 3730–3738. doi:10.1109/ICCV.2015.425. [6] E. Z. Liu, B. Haghgoo, A. S. Chen, A. Raghunathan, P. W. Koh, S. Sagawa, P. Liang, C. Finn, Just train twice: Improving group robustness without training group information, in: M. Meila, T. Zhang

Page 15 of 17

ProME: Group Robustness without Training Group Labels (Eds.), Proceedings of the 38th International Conference on Machine Learning, Vol. 139 of Proceedings of Machine Learning Research, PMLR, 2021, pp. 6781–6792. URL https://proceedings.mlr.press/v139/liu21f.html [7] M. Pezeshki, D. Bouchacourt, M. Ibrahim, N. Ballas, P. Vincent, D. Lopez-Paz, Discovering environments with XRM, in: R. Salakhutdinov, Z. Kolter, K. Heller, A. Weller, N. Oliver, J. Scarlett, F. Berkenkamp (Eds.), Proceedings of the 41st International Conference on Machine Learning, Vol. 235 of Proceedings of Machine Learning Research, PMLR, 2024, pp. 40551–40569. URL https://proceedings.mlr.press/v235/pezeshki24a.html [8] N. Sohoni, J. Dunnmon, G. Angus, A. Gu, C. Ré, No subclass left behind: Fine-grained robustness in coarse-grained classification problems, in: H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, H. Lin (Eds.), Advances in Neural Information Processing Systems, Vol. 33, Curran Associates, Inc., 2020, pp. 19339–19352. URL https://proceedings.neurips.cc/paper_files/paper/2020/file/ e0688d13958a19e087e123148555e4b4-Paper.pdf

[9] E. Creager, J.-H. Jacobsen, R. Zemel, Environment inference for invariant learning, in: M. Meila, T. Zhang (Eds.), Proceedings of the 38th International Conference on Machine Learning, Vol. 139 of Proceedings of Machine Learning Research, PMLR, 2021, pp. 2189– 2200. URL https://proceedings.mlr.press/v139/creager21a.html [10] Y. Han, D. Zou, Improving group robustness on spurious correlation requires preciser group inference, in: R. Salakhutdinov, Z. Kolter, K. Heller, A. Weller, N. Oliver, J. Scarlett, F. Berkenkamp (Eds.), Proceedings of the 41st International Conference on Machine Learning, Vol. 235 of Proceedings of Machine Learning Research, PMLR, 2024, pp. 17480–17504. URL https://proceedings.mlr.press/v235/han24g.html [11] M. Zhang, N. S. Sohoni, H. R. Zhang, C. Finn, C. Re, Correct-ncontrast: a contrastive approach for improving robustness to spurious correlations, in: K. Chaudhuri, S. Jegelka, L. Song, C. Szepesvari, G. Niu, S. Sabato (Eds.), Proceedings of the 39th International Conference on Machine Learning, Vol. 162 of Proceedings of Machine Learning Research, PMLR, 2022, pp. 26484–26516. URL https://proceedings.mlr.press/v162/zhang22z.html [12] S. Qiu, A. Potapczynski, P. Izmailov, A. G. Wilson, Simple and fast group robustness by automatic feature reweighting, in: A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, J. Scarlett (Eds.), Proceedings of the 40th International Conference on Machine Learning, Vol. 202 of Proceedings of Machine Learning Research, PMLR, 2023, pp. 28448–28467. URL https://proceedings.mlr.press/v202/qiu23c.html [13] Y. Liao, Q. Wu, Y. Wu, X. Yan, Decorr: Environment partitioning for invariant learning and ood generalization, Neural Networks 199 (2026) 108727. doi:https://doi.org/10.1016/j.neunet.2026.108727. URL https://www.sciencedirect.com/science/article/pii/ S0893608026001899

[14] S. Wang, M. Sun, Q. Huang, Y. Wang, Environment promoted invariant information learning for graph out-of-distribution generalization, Artificial Intelligence 355 (2026) 104522. doi:https://doi.org/10.1016/j.artint.2026.104522. URL https://www.sciencedirect.com/science/article/pii/ S0004370226000482

[15] P. Kirichenko, P. Izmailov, A. G. Wilson, Last layer re-training is sufficient for robustness to spurious correlations, in: ICML 2022: Workshop on Spurious Correlations, Invariance and Stability, 2022. URL https://openreview.net/forum?id=THOOBy1uWVH [16] T. LaBonte, V. Muthukumar, A. Kumar, Towards last-layer retraining for group robustness with fewer annotations, in: A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, S. Levine (Eds.), Advances in Neural Information Processing Systems, Vol. 36, Curran Associates, Inc., 2023, pp. 11552–11579. doi:10.52202/075280-0509. URL https://proceedings.neurips.cc/paper_files/paper/2023/file/ 265bee74aee86df77e8e36d25e786ab5-Paper-Conference.pdf

[17] J. C. Hill, T. LaBonte, X. Zhang, V. Muthukumar, On the unreasonable effectiveness of last-layer retraining, Transactions on Machine Learning Research (2026). URL https://openreview.net/forum?id=h81ztbrkFb [18] P. W. Koh, S. Sagawa, H. Marklund, S. M. Xie, M. Zhang, A. Balsubramani, W. Hu, M. Yasunaga, R. L. Phillips, I. Gao, T. Lee, E. David, I. Stavness, W. Guo, B. Earnshaw, I. Haque, S. M. Beery, J. Leskovec, A. Kundaje, E. Pierson, S. Levine, C. Finn, P. Liang, Wilds: A benchmark of in-the-wild distribution shifts, in: M. Meila, T. Zhang (Eds.), Proceedings of the 38th International Conference on Machine Learning, Vol. 139 of Proceedings of Machine Learning Research, PMLR, 2021, pp. 5637–5664. URL https://proceedings.mlr.press/v139/koh21a.html [19] M. Arjovsky, L. Bottou, I. Gulrajani, D. Lopez-Paz, Invariant risk minimization, arXiv preprint arXiv:1907.02893 (2019). doi:10. 48550/arXiv.1907.02893. URL https://doi.org/10.48550/arXiv.1907.02893 [20] B. Y. Idrissi, M. Arjovsky, M. Pezeshki, D. Lopez-Paz, Simple data balancing achieves competitive worst-group-accuracy, in: B. Schölkopf, C. Uhler, K. Zhang (Eds.), Proceedings of the First Conference on Causal Learning and Reasoning, Vol. 177 of Proceedings of Machine Learning Research, PMLR, 2022, pp. 336–351. URL https://proceedings.mlr.press/v177/idrissi22a.html [21] J. Nam, H. Cha, S. Ahn, J. Lee, J. Shin, Learning from failure: Debiasing classifier from biased classifier, in: H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, H. Lin (Eds.), Advances in Neural Information Processing Systems, Vol. 33, Curran Associates, Inc., 2020, pp. 20673–20684. URL https://proceedings.neurips.cc/paper_files/paper/2020/file/ eddc3427c5d77843c2253f1e799fe933-Paper.pdf

[22] G. Li, J. Liu, W. Hu, Bias amplification enhances minority group performance, Transactions on Machine Learning Research (2024). URL https://openreview.net/forum?id=75OwvzZZBT [23] G. Zheng, W. Ye, A. Zhang, Learning robust classifiers with selfguided spurious correlation mitigation, in: Proceedings of the ThirtyThird International Joint Conference on Artificial Intelligence, IJCAI ’24, 2024. doi:10.24963/ijcai.2024/619. URL https://doi.org/10.24963/ijcai.2024/619 [24] H. Han, S. Kim, H. Joo, S. Hong, J. Lee, Mitigating spurious correlations via disagreement probability, in: A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, C. Zhang (Eds.), Advances in Neural Information Processing Systems, Vol. 37, Curran Associates, Inc., 2024, pp. 74363–74382. doi:10.52202/079017-2366. URL https://proceedings.neurips.cc/paper_files/paper/2024/file/ 879c5890a9d2ecdcb590c9674cda4a59-Paper-Conference.pdf

[25] C. Tsirigotis, J. Monteiro, P. Rodriguez, D. Vazquez, A. Courville, Group robust classification without any group information, in: A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, S. Levine (Eds.), Advances in Neural Information Processing Systems, Vol. 36, Curran Associates, Inc., 2023, pp. 56553–56575. doi:10.52202/075280-2468. URL https://proceedings.neurips.cc/paper_files/paper/2023/file/ b0d9ceb3d11d013e55da201d2a2c07b2-Paper-Conference.pdf

[26] P. Q. Le, J. Schlötterer, C. Seifert, Out of spuriousity: Improving robustness to spurious correlations without group annotations, Transactions on Machine Learning Research (2025). URL https://openreview.net/forum?id=EEeVYfXor5 [27] J. P. Cohen, L. Blankemeier, A. S. Chaudhari, Identifying spurious correlations using counterfactual alignment, Transactions on Machine Learning Research (2025). URL https://openreview.net/forum?id=Utjw2z1ale [28] Z. Chen, M. Zhang, S. Cui, H. Li, G. Niu, M. Gong, C. Zhang, K. Zhang, Neural collapse inspired feature alignment for out-ofdistribution generalization, in: A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, C. Zhang (Eds.), Advances in Neural Information Processing Systems, Vol. 37, Curran Associates, Inc., 2024, pp. 93671–93689. doi:10.52202/079017-2970. URL https://proceedings.neurips.cc/paper_files/paper/2024/file/ aa56c74513a5e35768a11f4e82dd7ffb-Paper-Conference.pdf

Q. Wang et al.: Preprint submitted to Elsevier

Page 16 of 17

ProME: Group Robustness without Training Group Labels [29] N. Stromberg, R. Ayyagari, M. Welfert, S. Koyejo, R. Nock, L. Sankar, Robustness to subpopulation shift with domain label noise via regularized annotation of domains, Transactions on Machine Learning Research 2024, publisher Copyright: © 2024, Transactions on Machine Learning Research. All rights reserved. (2024). [30] J. Peters, P. Bühlmann, N. Meinshausen, Causal inference by using invariant prediction: Identification and confidence intervals, Journal of the Royal Statistical Society Series B: Statistical Methodology 78 (5) (2016) 947–1012. arXiv:https://academic.oup.com/ jrsssb/article-pdf/78/5/947/49235444/jrsssb_78_5_947.pdf, doi:10. 1111/rssb.12167. URL https://doi.org/10.1111/rssb.12167 [31] M. Rojas-Carulla, B. Schölkopf, R. Turner, J. Peters, Invariant models for causal transfer learning, Journal of Machine Learning Research 19 (36) (2018) 1–34. URL http://jmlr.org/papers/v19/16-432.html [32] K. Ahuja, K. Shanmugam, K. Varshney, A. Dhurandhar, Invariant risk minimization games, in: H. D. III, A. Singh (Eds.), Proceedings of the 37th International Conference on Machine Learning, Vol. 119 of Proceedings of Machine Learning Research, PMLR, 2020, pp. 145– 155. URL https://proceedings.mlr.press/v119/ahuja20a.html [33] J. Liu, Z. Hu, P. Cui, B. Li, Z. Shen, Integrated latent heterogeneity and invariance learning in kernel space, in: M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, J. W. Vaughan (Eds.), Advances in Neural Information Processing Systems, Vol. 34, Curran Associates, Inc., 2021, pp. 21720–21731. URL https://proceedings.neurips.cc/paper_files/paper/2021/file/ b59a51a3c0bf9c5228fde841714f523a-Paper.pdf

[34] Y. Lin, S. Zhu, L. Tan, P. Cui, Zin: When and how to learn invariance without environment partition?, in: S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, A. Oh (Eds.), Advances in Neural Information Processing Systems, Vol. 35, Curran Associates, Inc., 2022, pp. 24529–24542. doi:10.52202/068431-1781. URL https://proceedings.neurips.cc/paper_files/paper/2022/file/ 9b77f07301b1ef1fe810aae96c12cb7b-Paper-Conference.pdf

[35] E. Rosenfeld, P. Ravikumar, A. Risteski, The risks of invariant risk minimization, International Conference on Learning Representations 9. URL https://par.nsf.gov/biblio/10222689 [36] P. Kamath, A. Tangella, D. Sutherland, N. Srebro, Does invariant risk minimization capture invariance?, in: A. Banerjee, K. Fukumizu (Eds.), Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, Vol. 130 of Proceedings of Machine Learning Research, PMLR, 2021, pp. 4069–4077. URL https://proceedings.mlr.press/v130/kamath21a.html [37] J. Snell, K. Swersky, R. Zemel, Prototypical networks for few-shot learning, in: Proceedings of the 31st International Conference on Neural Information Processing Systems, NIPS’17, Curran Associates Inc., Red Hook, NY, USA, 2017, p. 4080–4090. [38] V. Papyan, X. Y. Han, D. L. Donoho, Prevalence of neural collapse during the terminal phase of deep learning training, Proceedings of the National Academy of Sciences 117 (40) (2020) 24652–24663. arXiv:https://www.pnas.org/doi/pdf/10.1073/pnas. 2015509117, doi:10.1073/pnas.2015509117. URL https://www.pnas.org/doi/abs/10.1073/pnas.2015509117 [39] M. N. N. To, P. F. R. Wilson, V. Nguyen, M. Harmanani, M. Cooper, F. Fooladgar, P. Abolmaesumi, P. Mousavi, R. G. Krishnan, Diverse prototypical ensembles improve robustness to subpopulation shift, in: Proceedings of the 42nd International Conference on Machine Learning, ICML’25, JMLR.org, 2025. [40] P. Izmailov, P. Kirichenko, N. Gruver, A. G. Wilson, On feature learning in the presence of spurious correlations, in: S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, A. Oh (Eds.), Advances in Neural Information Processing Systems, Vol. 35, Curran Associates, Inc., 2022, pp. 38516–38532. doi:10.52202/068431-2791. URL https://proceedings.neurips.cc/paper_files/paper/2022/file/

[41] T. LaBonte, J. C. Hill, X. Zhang, V. Muthukumar, A. Kumar, The group robustness is in the details: Revisiting finetuning under spurious correlations, in: A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, C. Zhang (Eds.), Advances in Neural Information Processing Systems, Vol. 37, Curran Associates, Inc., 2024, pp. 121598–121629. doi:10.52202/079017-3865. URL https://proceedings.neurips.cc/paper_files/paper/2024/file/ dc4db2ff2c1aefce3b594f821ea82fe6-Paper-Conference.pdf

[42] N. Stromberg, R. Ayyagari, M. Welfert, S. Koyejo, R. Nock, L. Sankar, For robust worst-group accuracy, ignore group annotations, Transactions on Machine Learning Research (2024). URL https://openreview.net/forum?id=l8E68fD6yp [43] T. G. J. Rudner, Y. Shi Zhang, A. G. Wilson, J. Kempe, Mind the GAP: Improving robustness to subpopulation shifts with group-aware priors, in: S. Dasgupta, S. Mandt, Y. Li (Eds.), Proceedings of The 27th International Conference on Artificial Intelligence and Statistics, Vol. 238 of Proceedings of Machine Learning Research, PMLR, 2024, pp. 127–135. URL https://proceedings.mlr.press/v238/rudner24a.html [44] N. Stromberg, R. Ayyagari, S. Koyejo, R. Nock, L. Sankar, Enhancing robustness of last layer two-stage fair model corrections, in: A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, C. Zhang (Eds.), Advances in Neural Information Processing Systems, Vol. 37, Curran Associates, Inc., 2024, pp. 139380–139401. doi:10.52202/079017-4424. URL https://proceedings.neurips.cc/paper_files/paper/2024/file/ fba9133d7fa896ab9414ddd1a6b1ecbf-Paper-Conference.pdf

[45] D. Krueger, E. Caballero, J.-H. Jacobsen, A. Zhang, J. Binas, D. Zhang, R. L. Priol, A. Courville, Out-of-distribution generalization via risk extrapolation (rex), in: M. Meila, T. Zhang (Eds.), Proceedings of the 38th International Conference on Machine Learning, Vol. 139 of Proceedings of Machine Learning Research, PMLR, 2021, pp. 5815–5826. URL https://proceedings.mlr.press/v139/krueger21a.html [46] D. Borkan, L. Dixon, J. Sorensen, N. Thain, L. Vasserman, Nuanced metrics for measuring unintended bias with real data for text classification, in: Companion Proceedings of The 2019 World Wide Web Conference, WWW ’19, Association for Computing Machinery, New York, NY, USA, 2019, p. 491–500. doi:10.1145/3308560.3317593. URL https://doi.org/10.1145/3308560.3317593 [47] S. Sagawa, A. Raghunathan, P. W. Koh, P. Liang, An investigation of why overparameterization exacerbates spurious correlations, in: H. D. III, A. Singh (Eds.), Proceedings of the 37th International Conference on Machine Learning, Vol. 119 of Proceedings of Machine Learning Research, PMLR, 2020, pp. 8346–8356. URL https://proceedings.mlr.press/v119/sagawa20a.html [48] J. Nam, J. Kim, J. Lee, J. Shin, Spread spurious attribute: Improving worst-group accuracy with spurious attribute estimation, in: International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id=_F9xpOrqyX9 [49] X. Zhou, Y. Lin, R. Pi, W. Zhang, R. Xu, P. Cui, T. Zhang, Model agnostic sample reweighting for out-of-distribution learning, in: K. Chaudhuri, S. Jegelka, L. Song, C. Szepesvari, G. Niu, S. Sabato (Eds.), Proceedings of the 39th International Conference on Machine Learning, Vol. 162 of Proceedings of Machine Learning Research, PMLR, 2022, pp. 27203–27221. URL https://proceedings.mlr.press/v162/zhou22d.html [50] R. Qiao, Z. Wu, J. Wang, P. W. Koh, B. K. H. Low, Group-robust sample reweighting for subpopulation shifts via influence functions, in: Y. Yue, A. Garg, N. Peng, F. Sha, R. Yu (Eds.), International Conference on Learning Representations, Vol. 2025, 2025, pp. 36959–36980. URL https://proceedings.iclr.cc/paper_files/paper/2025/file/ 5b6c1a77153c9d11861e26f92780cf2f-Paper-Conference.pdf

fb64a552feda3d981dbe43527a80a07e-Paper-Conference.pdf

Q. Wang et al.: Preprint submitted to Elsevier

Page 17 of 17

Related documents

• Dual-Channel Robust Group-Relative Policy Optimization via Advantage and Sequence-Weight Estimation#1132145
• Learning-Augmented Robust Algorithmic Recourse#173226
• Efficient Robust Learning at the Information-Theoretic Limit#1011384
• Local Evidence and Geometric Readout Repair in Trained GNNs#1063303
• Duality of causal distributionally robust optimization#617745