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

ProME: Prototype-Margin Environments with Repair-Aware Selection

for Group-Robust Learning

Summary

This paper introduces ProME (Prototype-Margin Environments), a two-stage framework for group-robust learning that operates without training-group labels. The authors formulate the problem as endogenous environments with repair-aware selection (ERAS) (Definition 1), which aligns both environment construction and model selection with the deployed predictor.

Problem formulation. Group-robust learning aims to maintain accuracy on rare subpopulations when training-group labels are unavailable. Existing methods typically use a two-stage pipeline where a reference model infers environments (by clustering features or identifying misclassified examples) and a separate predictor is trained on fixed assignments. This creates an environment–representation mismatch that weakens invariant learning. Additionally, model selection is often performed before classifier repair, which can reject checkpoints that would have performed well after repair. The ERAS formulation addresses both misalignments: environment construction is endogenous (drawn from the same trajectory it regularizes), and model selection is repair-aware (ranking candidates by worst-group accuracy after classifier repair).

ProME framework. Stage 1 uses a cosine-prototype classifier where each class is represented by a normalized prototype (Eq. 9). The prototype margin (Eq. 11) measures the difference in cosine support between the observed class and the strongest competing class. ProME splits these margins at their median to create two approximately balanced environments (Eq. 12). The training objective (Eq. 14) combines average environment risk with IRMv1 and REx penalties. Prototypes are refreshed periodically, and the partition can be updated (ProME-Refresh) or kept fixed. Stage 2 freezes each retained encoder checkpoint, fits a group-balanced linear head on group-annotated validation data (Eq. 17), and selects the encoder–head pair with the highest validation worst-group accuracy (Eq. 18). The selected pair is deployed as a single model.

Theoretical contributions. The paper provides four main theoretical results:

• Lemma 1 shows that the prototype margin and the prediction logit share one linear direction.

• Proposition 1 establishes that low prototype margins mark shortcut-conflicting examples when the causal and spurious components are sufficiently distinguishable (Eq. 22).

• Proposition 2 bounds the fraction of reassigned examples when the representation changes, showing partition stability (Eq. 24).

• Proposition 3 bounds the worst inferred-environment risk by the mean and variance of environment risks (Eq. 25). Corollary 1 shows this bound transfers to oracle groups under an explicit total-variation alignment condition (Eq. 26).

Experimental results. ProME is evaluated on Waterbirds, CelebA, CivilComments, and ColoredMNIST. Key findings include:

• ProME achieves the highest average worst-group accuracy (87.0%) among methods with the same group-label access, compared with 83.9% for the best baseline (GSR). It ranks first on Waterbirds (93.1%) and CivilComments (78.7%), and second on CelebA (89.3%).

• Prototype margins concentrate shortcut conflicts: on Waterbirds, conflicting examples constitute 0.68 of the low-margin environment and 0.01 of the high-margin environment (Fig. 3).

• Trajectory-derived margins improve pre-repair WGA from 68.95% to 78.85% compared with a frozen reference encoder (Table 4).

• Classifier repair narrows performance differences among Stage 1 variants: pre-repair WGA spans 12.20 points across variants, but matched repair reduces the final spread to 1.25 points (Table 4).

• Multi-candidate selection improves post-repair WGA on CelebA: single-checkpoint reaches 86.94%, while milestone and random pools reach 89.31% and 89.87% respectively (Fig. 5c).

• Classifier repair reduces the posterior-alignment gap on ColoredMNIST, with ProME having the smallest gap after repair (Fig. 6).

• Validation-group labels for classifier repair are most valuable when training minority support is limited: the gap between train-DFR and val-DFR is 18.68 points on Waterbirds (smallest group: 56 examples) but only 1.83 points on CelebA (smallest group: 1,387 examples) (Table 5).

• Results are stable across seed budgets: standard deviations of 0.62, 0.51, and 0.30 WGA points for Waterbirds, CelebA, and CivilComments across ten seeds (Fig. 7).

Conclusion. ProME offers deployment alignment as a practical design principle for group-robust learning. It requires no training-group labels, is backbone-agnostic, and deploys a single encoder–head pair. The authors hope this work facilitates future research across diverse forms of subpopulation shift.

Improvements to AI systems:

1. Endogenous environment construction for invariant learning. Instead of using a fixed, pre-trained reference model to define training environments, the AI system can dynamically construct environments from its own current representation during training. This eliminates the environment–representation mismatch, allowing the model to learn invariant features that are aligned with its final deployed representation, rather than features aligned with an outdated reference.

2. Repair-aware model selection. The AI system can select its final checkpoint based on performance after applying a classifier-repair step (e.g., fitting a balanced head on validation group labels), rather than on pre-repair performance. This prevents discarding high-potential encoders that would have excelled post-repair, improving worst-group accuracy without requiring additional training compute.

3. Prototype-margin-based environment splitting. The system can use the cosine margin between the observed class prototype and the strongest competing class prototype to partition training data into low- and high-margin environments. This provides a lightweight, label-free heuristic that concentrates shortcut-conflicting examples into the low-margin environment, enabling targeted invariant regularization without needing group annotations.

4. Trajectory-derived environment partitions with refresh. The system can periodically refresh its environment partition during training, using its evolving prototype margins. This allows the environment definition to track the model’s changing representation, improving robustness to spurious correlations compared to static partitions from a frozen encoder.

5. Multi-candidate encoder–head selection. The system can retain multiple encoder checkpoints from different training stages, freeze them, fit a group-balanced linear head on each, and select the pair with the highest validation worst-group accuracy. This improves final robustness by leveraging diversity across the training trajectory, especially when a single checkpoint may overfit or underfit.

6. Group-balanced linear head fitting for classifier repair. After training the main encoder, the system can replace the final classification layer with a linear head trained on a small set of group-annotated validation data, using class-balanced sampling. This repairs the classifier’s bias toward majority groups without retraining the encoder, and is most beneficial when training minority support is scarce.

7. Risk-variance-aware invariant regularization. The system can combine average environment risk with IRMv1 and REx penalties, where the theoretical bound (Proposition 3) ensures that worst-environment risk is controlled by the mean and variance of environment risks. This provides a principled objective that directly targets worst-group performance, even when environment labels are inferred.

8. Stable environment partitioning under representation shift. The system can rely on prototype margins that are provably stable (Proposition 2), bounding the fraction of examples reassigned when the representation changes. This ensures that periodic refreshes do not cause chaotic environment flips, leading to more stable training and reproducible results across random seeds.

What the improved AI system can do:

• Achieve higher worst-group accuracy on datasets with spurious correlations (e.g., Waterbirds, CelebA, CivilComments) without requiring any training-group labels, matching or exceeding methods that use such labels.

• Automatically identify and prioritize shortcut-conflicting examples during training, focusing invariant learning on the most challenging subpopulations.

• Select its final model in a way that is robust to classifier bias, leading to better performance on rare groups after deployment.

• Adapt its notion of environment over time, tracking its own learning progress and avoiding misalignment with a static reference model.

• Provide stable, reproducible robustness improvements across different random seeds and backbones, making it practical for real-world deployment where group labels are unavailable.