Rank-Constrained Adaptation for Reliable Real-World Performance

Abinitha Gourabathina; Hyewon Jeong; Teya Bergamaschi; Marzyeh Ghassemi; Collin Stultz

cs.LGFeb 6, 2026

Rank-Constrained Adaptation for Reliable Real-World Performance

Authors: Abinitha Gourabathina, Hyewon Jeong, Teya Bergamaschi, Marzyeh Ghassemi, Collin Stultz

Organizations: Department of Electrical Engineering & Computer Science, Massachusetts Institute of Technology, Cambridge, MA, USA · Harvard School of Medicine, Harvard, Boston, MA, USA.

Abstract

Deep learning models trained to optimize average accuracy often exhibit systematic failures on particular subpopulations. In real-world settings like healthcare, the subpopulations most affected by such disparities are frequently unlabeled, partially observed, or not known in advance. Existing group-robust methods typically assume prior knowledge of the relevant subgroups, using group annotations for training, validation, or model selection. We propose Misclassification Aware Rank-Limited Adaptation (MARLA), a parameter-efficient method for improving worst group performance without explicit subgroup annotations. MARLA leverages an ERM-trained model by calculating the model's misclassification probability scores on a held-out adaptation set to identify a low-dimensional subspace where errors concentrate. We then learn a rank-restricted additive correction to the classifier logits within that subspace. Across seven real-world datasets, we evaluate group robustness under three settings: no knowledge of subgroup relevance, partial knowledge of subgroup relevance, and full knowledge of subgroup relevance. MARLA improves worst-group performance while remaining fast and parameter-efficient, with data-guided hyperparameter selection.

Figures & tables

Figure 1 : MARLA. ( Left ) ERM can learn a decision boundary that performs well on average while systematically misclassifying a subpopulation. ( Middle ) MARLA assigns larger weights to examples where the model gives low probability to the true label, and uses these weights to identify important directions in representation space. ( Right ) MARLA learns a rank-restricted additive correction to the classifier logits along these directions, improving classification for affected samples without unrestricted classifier retraining. Colors denote class labels, while marker shapes denote attribute groups.

Algorithm	Group Info	GRACE		MIMIC-IV
Algorithm	(Train/Val)	WGA	Avg Acc	WGA	Avg Acc
ERM*	✗/✗	$93.3\scriptstyle{\pm 0.7}$	$95.1\scriptstyle{\pm 0.4}$	$72.1\scriptstyle{\pm 0.0}$	$80.0\scriptstyle{\pm 0.0}$
IRM*	✗/✗	$89.9\scriptstyle{\pm 1.3}$	$93.7\scriptstyle{\pm 0.6}$	$64.7\scriptstyle{\pm 13.2}$	$75.9\scriptstyle{\pm 6.5}$
CVaRDRO*	✗/✗	$88.7\scriptstyle{\pm 2.5}$	$92.1\scriptstyle{\pm 1.4}$	$38.3\scriptstyle{\pm 11.4}$	$54.8\scriptstyle{\pm 6.9}$
JTT*	✗/✗	$67.2\scriptstyle{\pm 3.9}$	$74.7\scriptstyle{\pm 6.3}$	$72.1\scriptstyle{\pm 0.0}$	$79.2\scriptstyle{\pm 0.4}$
LfF*	✗/✗	$91.3\scriptstyle{\pm 1.0}$	$94.1\scriptstyle{\pm 0.9}$	$72.6\scriptstyle{\pm 1.5}$	$80.0\scriptstyle{\pm 0.2}$

Table 1 : Worst-group accuracy (WGA) and average accuracy on test set with no group information. Rows with asterisks (*) are reproduced or evaluated in our implementation; other results are reported from original papers. Best and second-best values are bold and underlined , respectively. “-” indicates results not reported. All results are the mean

\pm

std (in %) averaged over three runs.

\dagger

GSR is not supervision-matched in this setting because it uses group-labeled validation examples.

Algorithm	Group Info	GRACE				MIMIC-IV
	Known Attribute	Age		Sex		Race		Gender
	(Train/Val)	WGA	Avg Acc	WGA	Avg Acc	WGA	Avg Acc	WGA	Avg Acc
ERM	✗/★	$91.1\scriptstyle{\pm 1.1}$	$\underline{95.1}\scriptstyle{\pm 0.4}$	$93.3\scriptstyle{\pm 0.7}$	$95.1\scriptstyle{\pm 0.4}$	$72.1\scriptstyle{\pm 0.3}$	$80.0\scriptstyle{\pm 0.2}$	$72.1\scriptstyle{\pm 0.3}$	$80.0\scriptstyle{\pm 0.2}$
IRM	✗/★	$44.8\scriptstyle{\pm 9.6}$	$70.3\scriptstyle{\pm 6.1}$	$50.2\scriptstyle{\pm 8.1}$	$73.6\scriptstyle{\pm 5.4}$	$72.4\scriptstyle{\pm 0.5}$	$80.0\scriptstyle{\pm 0.3}$	$71.8\scriptstyle{\pm 0.4}$	$80.2\scriptstyle{\pm 0.2}$
CVaRDRO	✗/★	$47.5\scriptstyle{\pm 11.2}$	$72.0\scriptstyle{\pm 5.8}$	$52.6\scriptstyle{\pm 9.3}$	$75.4\scriptstyle{\pm 4.9}$	$76.3\scriptstyle{\pm 0.7}$	$\underline{80.3}\scriptstyle{\pm 0.1}$	$76.4\scriptstyle{\pm 0.6}$	$\underline{80.3}\scriptstyle{\pm 0.1}$
JTT	✗/★	$86.4\scriptstyle{\pm 4.3}$	$92.6\scriptstyle{\pm 1.1}$	$90.2\scriptstyle{\pm 3.7}$	$93.7\scriptstyle{\pm 1.2}$	$72.0\scriptstyle{\pm 0.4}$	$80.1\scriptstyle{\pm 0.2}$	$72.2\scriptstyle{\pm 0.3}$	$79.8\scriptstyle{\pm 0.4}$

Table 2: Worst-group accuracy (WGA) and average accuracy on the test set under partial information of group relevance .

Algorithm	Group Info	GRACE		MIMIC-IV
Algorithm	(Train/Val)	WGA	Avg Acc	WGA	Avg Acc
ERM*	✗/★	$\underline{93.3\scriptstyle{\pm 0.7}}$	$\underline{95.1\scriptstyle{\pm 0.4}}$	$72.1\scriptstyle{\pm 0.3}$	$80.0\scriptstyle{\pm 0.2}$
CVaRDRO*	✗/★	$88.9\scriptstyle{\pm 2.3}$	$93.5\scriptstyle{\pm 1.4}$	$\underline{76.2\scriptstyle{\pm 0.7}}$	$80.0\scriptstyle{\pm 0.6}$
JTT*	✗/★	$90.2\scriptstyle{\pm 3.7}$	$93.7\scriptstyle{\pm 1.2}$	$71.9\scriptstyle{\pm 0.4}$	$79.9\scriptstyle{\pm 0.3}$
LfF*	✗/★	$87.5\scriptstyle{\pm 2.8}$	$92.6\scriptstyle{\pm 1.6}$	$72.4\scriptstyle{\pm 0.5}$	$80.1\scriptstyle{\pm 0.2}$
Mixup*	✗/★	$88.3\scriptstyle{\pm 2.5}$	$92.9\scriptstyle{\pm 1.5}$	$71.8\scriptstyle{\pm 0.4}$	$79.8\scriptstyle{\pm 0.3}$

Table 3 : Worst-group accuracy (WGA) and average accuracy on the test set with complete information of group relevance. Rows with asterisks (*) and CheXpert results reproduce baselines. Group Info (Train/Val): ✗ = no group info used, ★ = used for tuning/early stopping. Best and second-best values within each half are bold and underlined . “-” indicates results not reported. All results are mean

\pm

std (in %) over three runs.

Figure 2 : MARLA performance under variable rank conditions (a) The score-weighted covariance spectrum (a central component of MARLA) exhibits a sharp elbow exactly at true rank $K$ (solid curves), while PCA on the unweighted covariance matrix shows shallow non-specific decay (dashed curves). (b) Worst-group accuracy recovers only once the correction rank (the number of eigenvectors used in MARLA) reaches $K$ ; PCA-directed correction yields no improvement at any rank. (c) Seven benchmark datasets exhibit concentrated score-weighted covariance spectra.

	Waterbirds		MultiNLI		CelebA		CivilComments		CheXpert
Method	WGA	Avg Acc	WGA	Avg Acc	WGA	Avg Acc	WGA	Avg Acc	WGA	Avg Acc
MARLA (Low-rank)	$\bf 90.1_{\pm 0.1}$	$\bf 93.7_{\pm 0.7}$	$\bf 69.7_{\pm 0.1}$	$\bf 81.2_{\pm 0.2}$	$\bf 82.8_{\pm 0.5}$	$\bf 95.2_{\pm 0.1}$	$\bf 71.6_{\pm 0.7}$	$\bf 91.4_{\pm 0.4}$	$\bf 75.3_{\pm 0.4}$	$79.6_{\pm 0.2}$
Full-rank	$82.0_{\pm 4.5}$	$92.0_{\pm 2.0}$	$69.7_{\pm 1.7}$	$81.0_{\pm 0.5}$	$81.3_{\pm 0.4}$	$90.2_{\pm 0.1}$	$68.7_{\pm 1.3}$	$91.0_{\pm 0.6}$	$56.4_{\pm 9.1}$	$79.1_{\pm 0.8}$
ERM	$69.1_{\pm 4.7}$	$84.1_{\pm 1.7}$	$66.4_{\pm 2.3}$	$81.0_{\pm 0.3}$	$57.6_{\pm 0.8}$	$95.0_{\pm 0.1}$	$63.2_{\pm 1.2}$	$85.4_{\pm 0.2}$	$41.7_{\pm 3.4}$	$88.6_{\pm 0.7}$

Table 4: Ablation study of low-rank vs full-rank correction: Worst-group accuracy (WGA) and average accuracy (Avg Acc) in the no group-relevance information setting. ERM values match the corresponding no-group-information ERM results in Table 1 .

Figure 3 : ( Left ) Training time (in min.) across datasets. MARLA only incurs a small overhead compared to standard ERM and is substantially less expensive than recent methods. ( Right ) Average number of trainable parameters updated in the adaptation from ERM for comparable baselines in an independent run. $d$ = embedding dimension, $C$ = number of classes, $N$ = DPE number of prototypes per class, and $P_{\text{bb}}$ = total parameters of base model backbone.

Appendix figures & tables24 assets

Supplementary material from the paper’s appendix.

Appendix

$N$	$\rho$	ERM UGA (%)	GDRO UGA (%)	Harm (%)	MARLA UGA (%)
1	0.1	$77.5\pm 15.5$	$55.4\pm 7.2$	$22.1\pm 8.5$	$85.5\pm 11.0$
1	0.2	$84.9\pm 0.4$	$49.1\pm 9.0$	$35.8\pm 9.2$	$91.5\pm 7.8$
1	0.3	$85.1\pm 4.7$	$44.8\pm 4.8$	$40.3\pm 9.1$	$86.9\pm 7.6$
2	0.1	$58.3\pm 7.6$	$42.5\pm 5.9$	$15.9\pm 3.6$	$71.0\pm 19.0$
2	0.2	$82.3\pm 4.0$	$42.2\pm 11.2$	$40.1\pm 8.8$	$89.6\pm 4.8$
2	0.3	$75.6\pm 4.9$	$49.8\pm 9.1$	$25.8\pm 13.3$	$88.7\pm 1.4$

Appendix

Table 5: Unknown-group accuracy (UGA) across synthetic parameter sweeps.

N

denotes the number of known groups and

\rho

denotes the unknown group size ratio. UGA is reported for ERM and Group DRO. Harm is defined as the percentage-point decrease in UGA when using Group DRO relative to ERM.

Figure 4 : Performance as a function of the unknown-group size ratio. Left: unknown-group accuracy for ERM and Group DRO, averaged over the number of known groups. ERM accuracy improves as the unknown subgroup becomes larger, while Group DRO remains approximately constant or decreases. Right: harm, defined as ERM accuracy minus Group DRO accuracy, increases with the unknown-group size ratio.

Figure 5 : Performance on the unknown subgroup across parameter settings. The heatmaps show harm on the unknown subgroup (top left), ERM accuracy (top right), Group DRO accuracy (bottom left), and harm as a function of the unknown-group size ratio for different numbers of known groups (bottom right). Results are shown for 2-5 known groups and unknown-group size ratios from 0.1 to 0.5. Warmer colors indicate larger values.

Dataset	Examples
Image datasets
Waterbirds
CelebA
CheXpert
Text datasets
CivilComments	“Munchins looks like a munchins. The man who dont want to show his taxes, will tell you everything…” “The democratic party removed the filibuster to steamroll its agenda. Suck it up boys and girls.” “so you dont use oil? no gasoline? no plastic? man you ignorant losers are pathetic.”

Appendix

Table 6: Example inputs from datasets

OUTCOME
	Train (N=258,944)	Val (N=86,589)	Test (N=85,698)	Overall (N=431,231)
30-day readmission	52,580 (20.3%)	17,404 (20.1%)	17,102 (20.0%)	87,086 (20.2%)
In-hospital death	5,178 (2.0%)	1,704 (2.0%)	1,727 (2.0%)	8,609 (2.0%)
DEMOGRAPHICS
Age, years (mean, SD)	56.8 (19.0)	57.1 (19.1)	56.9 (19.0)	56.9 (19.1)
Female sex	134,251 (51.8%)	45,557 (52.6%)	45,182 (52.7%)	224,990 (52.2%)

Appendix

Table 7 : Cohort Characteristics of MIMIC-IV

OUTCOME
	Train (N=13,442)	Val (N=4,480)	Test (N=4,482)	Overall (N=22,404)
In-hospital death	472 (3.5%)	172 (3.8%)	173 (3.9%)	817 (3.6%)
DEMOGRAPHICS
Age, years (mean, SD)	65.9 (13.6)	65.7 (13.7)	65.9 (13.7)	65.9 (13.7)
Female sex	4,358 (32.4%)	1,464 (32.7%)	1,429 (31.9%)	7,251 (32.4%)
Male sex	9,084 (67.6%)	3,016 (67.3%)	3,053 (68.1%)	15,153 (67.6%)

Appendix

Table 8 : Cohort Characteristics of GRACE

Dataset	Architecture	Epochs	Batch Size	Learning Rate	Weight Decay	Scheduler	Optimizer
Waterbirds	ResNet-50 (ImageNet)	50	32	0.003	$1\times 10^{-4}$	Constant	SGD
CelebA	ResNet-50 (ImageNet)	20	100	0.003	$1\times 10^{-4}$	Cosine	SGD
CivilComments	BERT-base-uncased	3	32	$2\times 10^{-5}$	0.0	Constant	AdamW
MultiNLI	BERT-base-uncased	3	32	$2\times 10^{-5}$	0.0	Constant	AdamW
CheXpert	ResNet-50 (ImageNet)	13 ∗	108	0.001	$1\times 10^{-4}$	Default ∗∗	SGD
GRACE	MLP (6 layers, 256 hidden)	5,001 †	1024	$1\times 10^{-3}$	$1\times 10^{-5}$	Constant	SGD

Appendix

Table 9 : Base ERM Training Configuration for All Datasets

Dataset	Epochs	Gamma Values	Rank Values
Waterbirds	100	7.5, 8.5	4, 16, 32
CelebA	100	0.5, 0.75, 1.25, 1.5	2, 4, 12, 36, 52
CivilComments	50	500, 1000, 5000	1, 4, 5
MultiNLI	50	200, 500, 1000	1, 2, 3
CheXpert	50	0.5, 1, 1.5	50, 52, 54
GRACE	300	1.0, 5.0, 10.0, 20.0, 50.0	1, 2, 4, 8, 16

Appendix

Table 10 : MARLA Stage 2 Hyperparameter Configurations Across Datasets

Method	Search Type	Trials	Epochs	Learning Rate	Hyperparameters
MARLA	Grid	10	100	$2\times 10^{-2}$	$\gamma\in[100,1000]$ , $k\in[1,4,5,6,9]$ , $\lambda=0$
AFR (CC)	Grid	12	100	$3\times 10^{-3}$	$\gamma\in[0,0.01,0.1,1,3,10]$ , $\lambda\in[0.0,1\times 10^{-2}]$
AFR (CX)	Grid	32	100	$[10^{-4},10^{-3},10^{-2},5{\times}10^{-2}]$	$\gamma\in[1,2,4,6]$ , $\lambda\in[0.0,1\times 10^{-2}]$
DFR (CC)	Random	25	100	$1\times 10^{-2}$	$\text{dfr\_reg}\sim 10^{\text{Uniform}(-2.0,0.5)}$
DFR (CX)	Random	25	100	$1\times 10^{-2}$	$\text{dfr\_reg}\sim 10^{\text{Uniform}(-0.5,0.5)}$
CRT	None	1	100	$1\times 10^{-2}$	None

Appendix

Table 11 : Hyperparameter Sweep Configurations for Baseline Methods in for Partial Group Relevance Setting

Dataset	ERM	MARLA	AFR	JTT	DPE	GIC
Waterbirds	25.05	26.53	26.58	76.08	40.12	120.20
CelebA	158.36	167.03	169.27	479.71	254.00	767.24
CivilComments	145.95	164.89	166.68	441.85	257.51	757.72
MultiNLI	61.23	72.99	73.12	185.52	116.33	335.29

Appendix

Table 12 : End-to-end Runtime Comparison (minutes)

Dataset	Error-weight correlation	Worst-group amplification	Amplification relative to other groups	ERM error on worst group
GRACE	$0.96\pm 0.02$	$1.22\pm 0.07\times$	$1.68\pm 0.20\times$	$6.7\%\pm 0.7\%$
MIMIC-IV	$0.93\pm 0.04$	$4.1\pm 0.6\times$	$5.7\pm 0.9\times$	$27.9\%\pm 1.0\%$
Waterbirds	$0.99\pm 0.01$	$12.7\pm 3.9\times$	$18.4\pm 5.1\times$	$30.3\%\pm 6.7\%$
CivilComments	$0.94\pm 0.04$	$2.3\pm 0.4\times$	$3.2\pm 0.6\times$	$38.1\%\pm 8.7\%$

Appendix

Table 13 : Post-hoc relationship between ERM subgroup error and MARLA’s weight allocation. Quantities are computed per seed and reported as mean

\pm

standard deviation over three seeds.

Dataset	MARLA-attr	PCA-attr	Random-attr	MARLA-fix	PCA-fix	Random-fix
Waterbirds	0.44	0.21	0.01	0.52	0.42	0.02
CivilComments	0.55	0.68	0.02	0.75	0.22	0.02
GRACE	0.53	0.36	0.04	0.33	0.29	0.02
MIMIC-IV	0.52	0.31	0.03	0.64	0.05	0.01

Appendix

Table 14 : MSCC between candidate rank-

k

subspaces and post-hoc known-attribute or oracle error-fixing reference subspaces.

Dataset	MARLA WGA / Avg	PCA WGA / Avg	Random WGA / Avg	MARLA-PCA WGA	MARLA-Random WGA
Waterbirds	$\mathbf{90.1\pm 0.1\;/\;93.7}$	$76.6\pm 6.9\;/\;92.4$	$64.7\pm 8.7\;/\;90.1$	$\mathbf{+13.5}$	$\mathbf{+25.4}$
CivilComments	$\mathbf{71.6\pm 0.7\;/\;91.4}$	$54.2\pm 9.5\;/\;92.7$	$60.1\pm 0.4\;/\;92.5$	$\mathbf{+17.4}$	$\mathbf{+11.5}$
GRACE	$\mathbf{94.2\pm 1.6\;/\;95.8}$	$91.2\pm 0.8\;/\;95.1$	$91.1\pm 1.0\;/\;95.1$	$\mathbf{+3.0}$	$\mathbf{+3.1}$
MIMIC-IV	$\mathbf{78.2\pm 2.0\;/\;81.0}$	$76.9\pm 0.5\;/\;80.0$	$77.0\pm 0.5\;/\;80.1$	$\mathbf{+1.3}$	$\mathbf{+1.2}$

Appendix

Table 15 : Matched comparison of MARLA’s error-weighted basis with unweighted PCA and random orthonormal rank-

k

bases. Each entry is test WGA / average accuracy.

Figure 6 : Sensitivity of MARLA to correction rank $k$ . Test WGA remains stable across much of the tuning range. Curves show the mean and standard deviation over three seeds.

Figure 7 : MARLA Robustness to Gamma ( $\gamma$ ) Selection. Test worst-group accuracy (WGA) as a function of gamma for Waterbirds, CelebA, and CivilComments . For each gamma value, we select the best rank (by validation WGA) and report mean $\pm$ 1 standard deviation across seeds. The green dashed line and diamond marker indicate the empirically-observed optimal gamma for each dataset. The gray dashed line and shaded region show ERM performance (mean $\pm$ 1 std). MARLA demonstrates robustness to gamma selection, with performance remaining stable across a reasonable range around the optimal values.

Figure 8 : Differences in soft misclassification-score distributions by dataset. Image datasets have many examples with very low scores and a distinct high-score tail, whereas text datasets exhibit more compressed score distributions.

	Waterbirds		CivilComments
$\mathcal{D}_{\text{ERM}}/\mathcal{D}_{\text{MARLA{}}}$	WGA	Avg Acc	WGA	Avg Acc
90-10	$88.8\pm 0.3$	$92.3\pm 0.4$	$70.2\pm 0.3$	$91.0\pm 0.9$
80-20	$90.1\pm 0.1$	$93.7\pm 0.7$	$71.6\pm 0.7$	$91.4\pm 0.4$
70-30	$88.8\pm 0.3$	$92.3\pm 0.4$	$70.9\pm 0.1$	$91.1\pm 0.8$
60-40	$88.8\pm 0.3$	$92.3\pm 0.4$	$71.0\pm 0.8$	$91.1\pm 0.8$
ERM	$69.1\pm 4.7$	$84.1\pm 1.7$	$63.2\pm 1.2$	$85.4\pm 0.2$

Appendix

Table 16 : Performance comparison under varying

D_{\text{ERM}}/D_{\text{MARLA{}}}

ratios in the setting with no group information.

	Waterbirds		CivilComments
$\mathcal{D}_{\text{ERM}}/\mathcal{D}_{\text{MARLA{}}}$	WGA	Avg Acc	WGA	Avg Acc
90-10	$89.4\pm 0.3$	$92.8\pm 0.9$	$72.2\pm 0.7$	$91.4\pm 0.4$
80-20	$90.1\pm 0.1$	$92.8\pm 0.9$	$72.9\pm 0.2$	$90.9\pm 0.8$
70-30	$89.4\pm 0.3$	$92.3\pm 0.4$	$72.0\pm 0.3$	$91.3\pm 0.5$
60-40	$89.2\pm 0.3$	$92.5\pm 0.6$	$71.7\pm 0.2$	$91.3\pm 0.4$
ERM	$69.1\pm 4.7$	$84.1\pm 1.7$	$63.2\pm 1.2$	$85.4\pm 0.2$

Appendix

Table 17 : Performance comparison under varying

D_{\text{ERM}}/D_{\text{MARLA{}}}

ratios in the setting with complete group relevance information.

Method	AUROC	Macro AUPRC	Balanced acc	Sensitivity	Specificity	Age $\times$ sex WGA	Label-cond. demo WGA
ERM	$49.73\pm 11.26$	$51.75\pm 2.28$	$51.06\pm 1.50$	$2.27\pm 3.21$	$99.85\pm 0.21$	$92.52\pm 0.09$	$1.30\pm 1.84$
JTT	$49.73\pm 11.26$	$51.75\pm 2.28$	$51.06\pm 1.50$	$2.27\pm 3.21$	$99.85\pm 0.21$	$92.52\pm 0.09$	$1.30\pm 1.84$
IRM	$49.73\pm 11.26$	$51.75\pm 2.28$	$51.06\pm 1.50$	$2.27\pm 3.21$	$99.85\pm 0.21$	$92.52\pm 0.09$	$1.30\pm 1.84$
CVaRDRO	$49.77\pm 11.26$	$51.76\pm 2.28$	$51.06\pm 1.50$	$2.27\pm 3.21$	$99.85\pm 0.21$	$92.52\pm 0.09$	$1.30\pm 1.84$
LfF	$49.76\pm 11.25$	$51.76\pm 2.28$	$51.06\pm 1.49$	$2.27\pm 3.21$	$99.84\pm 0.23$	$92.52\pm 0.09$	$1.30\pm 1.84$
Mixup	$49.73\pm 11.26$	$51.75\pm 2.28$	$51.06\pm 1.50$	$2.27\pm 3.21$	$99.85\pm 0.21$	$92.52\pm 0.09$	$1.30\pm 1.84$

Appendix

Table 18 : Additional no-group-information metrics on GRACE.

Method	AUROC	Macro AUPRC	Balanced acc	Sensitivity	Specificity	Race $\times$ gender WGA	Label-cond. demo WGA
ERM	$63.38\pm 0.66$	$58.88\pm 0.51$	$59.76\pm 0.72$	$60.59\pm 1.92$	$58.93\pm 3.13$	$53.25\pm 1.52$	$36.15\pm 3.31$
JTT	$64.18\pm 0.14$	$59.43\pm 0.22$	$60.17\pm 0.42$	$63.59\pm 4.82$	$56.75\pm 5.66$	$53.22\pm 2.91$	$37.44\pm 1.81$
IRM	$58.80\pm 4.54$	$55.90\pm 2.73$	$56.09\pm 3.49$	$62.37\pm 1.78$	$49.82\pm 5.37$	$42.05\pm 0.49$	$30.41\pm 5.05$
CVaRDRO	$54.54\pm 2.73$	$53.18\pm 1.20$	$53.19\pm 2.33$	$53.76\pm 5.07$	$52.62\pm 3.34$	$44.46\pm 2.70$	$37.00\pm 1.90$
LfF	$58.45\pm 1.51$	$55.67\pm 1.12$	$55.62\pm 0.96$	$60.48\pm 3.13$	$50.76\pm 3.66$	$44.75\pm 0.94$	$36.50\pm 5.80$
Mixup	$63.09\pm 0.74$	$58.75\pm 0.63$	$59.48\pm 0.75$	$59.85\pm 1.96$	$59.11\pm 3.34$	$52.26\pm 1.02$	$35.13\pm 2.98$

Appendix

Table 19 : Additional no-group-information metrics on MIMIC-IV.

Dataset	Original $T=1$	Sharper $T=0.5$	Flatter $T=2$	Strongly flatter $T=4$	Calibrated $T_{\mathrm{cal}}$	Max absolute $\Delta$ WGA
Waterbirds	$90.1\pm 0.1\,/\,93.7\pm 0.7$	$90.7\pm 2.9\,/\,93.9\pm 2.8$	$90.6\pm 2.4\,/\,93.8\pm 3.2$	$90.5\pm 2.6\,/\,93.9\pm 3.0$	$90.5\pm 2.5\,/\,93.8\pm 3.0$	$0.6$ pp
CivilComments	$71.6\pm 0.7\,/\,91.4\pm 0.4$	$72.7\pm 2.2\,/\,91.2\pm 0.4$	$70.0\pm 0.7\,/\,91.5\pm 0.6$	$68.9\pm 1.5\,/\,91.6\pm 0.6$	$70.5\pm 0.6\,/\,91.5\pm 0.5$	$2.7$ pp
GRACE	$94.2\pm 1.6\,/\,95.8\pm 1.3$	$92.4\pm 3.1\,/\,94.7\pm 2.4$	$97.5\pm 0.6\,/\,98.3\pm 0.2$	$97.9\pm 1.0\,/\,98.5\pm 0.4$	$97.6\pm 0.7\,/\,98.3\pm 0.2$	$3.7$ pp
MIMIC-IV	$78.2\pm 2.0\,/\,81.0\pm 0.2$	$81.3\pm 2.5\,/\,82.0\pm 0.4$	$77.6\pm 4.4\,/\,80.5\pm 0.5$	$75.0\pm 7.0\,/\,79.0\pm 0.6$	$78.1\pm 3.6\,/\,81.0\pm 0.7$	$3.2$ pp

Appendix

Table 20 : Sensitivity of MARLA to temperature scaling used only for weight construction. Each entry is test WGA / average accuracy; the final column reports the largest absolute WGA change from

T=1

Group labels are not used for training
Algorithm	Group Info	GRACE		MIMIC-IV
Algorithm	(Train/Val)	WGA	Avg Acc	WGA	Avg Acc
ERM*	✗/★	$93.3\scriptstyle{\pm 0.7}$	$95.1\scriptstyle{\pm 0.4}$	$72.1\scriptstyle{\pm 0.3}$	$80.0\scriptstyle{\pm 0.2}$
CVaRDRO*	✗/★	$88.9\scriptstyle{\pm 2.3}$	$93.5\scriptstyle{\pm 1.4}$	$76.2\scriptstyle{\pm 0.7}$	$80.0\scriptstyle{\pm 0.6}$
JTT*	✗/★	$90.2\scriptstyle{\pm 3.7}$	$93.7\scriptstyle{\pm 1.2}$	$71.9\scriptstyle{\pm 0.4}$	$79.9\scriptstyle{\pm 0.3}$
LfF*	✗/★	$87.5\scriptstyle{\pm 2.8}$	$92.6\scriptstyle{\pm 1.6}$	$72.4\scriptstyle{\pm 0.5}$	$80.1\scriptstyle{\pm 0.2}$

Appendix

Table 21 : Worst-group accuracy (WGA) and average accuracy on the test set with complete information of group relevance. Baselines are grouped according to whether group labels are used during training. Rows marked with an asterisk (*) and the CheXpert results reproduce baselines from SubpopBench [ 97 ] ; other baseline results are taken from the corresponding papers. Group Info (Train/Val) indicates how group labels are used: ✗ = not used, ★ = used for hyperparameter tuning and early stopping, and ✓ = used for training. Best and second-best values within each half are bold and underlined , respectively. “-” indicates that a result was not reported. Our results are mean

\pm

standard deviation (in %) over three independent runs.

Group labels are not used for training
Algorithm	Group Info	Waterbirds		CelebA		CivilComments		MultiNLI		CheXpert
Algorithm	(Train/Val)	WGA	Avg Acc	WGA	Avg Acc	WGA	Avg Acc	WGA	Avg Acc	WGA	Avg Acc
ERM*	✗/★	$69.1\scriptstyle{\pm 4.7}$	$84.1\scriptstyle{\pm 1.7}$	$57.6\scriptstyle{\pm 0.8}$	$\underline{95.0}\scriptstyle{\pm 0.1}$	$63.2\scriptstyle{\pm 1.2}$	$85.4\scriptstyle{\pm 0.2}$	$69.5\scriptstyle{\pm 0.3}$	$80.9\scriptstyle{\pm 0.3}$	$41.7\scriptstyle{\pm 3.4}$	$\textbf{88.6}\scriptstyle{\pm 0.7}$
CRT*	✗/★	$76.3\scriptstyle{\pm 0.8}$	$89.2\scriptstyle{\pm 0.1}$	$70.4\scriptstyle{\pm}0.4$	$94.1\scriptstyle{\pm 0.1}$	$68.5\scriptstyle{\pm 0.0}$	$83.0\scriptstyle{\pm 0.0}$	$65.4\scriptstyle{\pm 0.1}$	$80.2\scriptstyle{\pm 0.0}$	$74.0\scriptstyle{\pm 0.2}$	$79.1\scriptstyle{\pm 0.1}$
ReWeightCRT*	✗/★	$76.3\scriptstyle{\pm 0.2}$	$89.4\scriptstyle{\pm 0.3}$	$71.1\scriptstyle{\pm}0.5$	$94.2\scriptstyle{\pm 0.1}$	$68.2\scriptstyle{\pm 0.4}$	$83.4\scriptstyle{\pm 0.0}$	$65.3\scriptstyle{\pm 0.1}$	$80.2\scriptstyle{\pm 0.0}$	$73.9\scriptstyle{\pm 0.2}$	$79.0\scriptstyle{\pm 0.1}$
JTT	✗/★	$86.7$	$93.3$	$81.1$	$88.0$	$69.3$	91.1	$\underline{72.6}$	$78.6$	$60.4\scriptstyle{\pm 4.8}$	$75.2\scriptstyle{\pm 0.8}$

Appendix

Table 30

		ERM		WGA (%)			Avg. Acc. (%)
Dataset	Model	WGA	Acc	MARLA	AFR	JTT	MARLA	AFR	JTT
Waterbirds	ResNet-101	$68.7{\pm}3.7$	$87.7{\pm}3.5$	$\mathbf{79.1{\pm}1.7}$	$72.5{\pm}1.0$	$69.3{\pm}1.3$	$86.0{\pm}0.8$	$86.1{\pm}0.7$	$86.6{\pm}2.1$
Waterbirds	ResNet-152	$72.2{\pm}0.7$	$88.6{\pm}2.6$	$\mathbf{81.5{\pm}0.8}$	$78.3{\pm}1.1$	$70.9{\pm}1.1$	$88.7{\pm}0.9$	$88.6{\pm}0.9$	$88.4{\pm}2.1$
CelebA	ResNet-101	$48.9{\pm}1.6$	$95.5{\pm}0.3$	$\mathbf{78.3{\pm}0.8}$	$73.6{\pm}4.2$	$61.4{\pm}6.9$	$90.1{\pm}1.2$	$92.2{\pm}1.1$	$94.6{\pm}0.8$
CelebA	ResNet-152	$47.4{\pm}3.0$	$95.9{\pm}0.1$	$\mathbf{73.1{\pm}3.8}$	$73.3{\pm}5.6$	$53.3{\pm}7.6$	$91.8{\pm}2.1$	$92.3{\pm}1.8$	$95.5{\pm}0.5$
CheXpert	ResNet-101	$28.8{\pm}3.2$	$88.9{\pm}0.5$	$\mathbf{76.3{\pm}1.8}$	$72.7{\pm}2.2$	$73.9{\pm}1.4$	$76.1{\pm}0.9$	$75.4{\pm}1.1$	$77.6{\pm}0.7$

Appendix

Table 22 : Performance of MARLA and baselines with larger backbones. WGA denotes worst-group accuracy. Results are mean

\pm

standard deviation over three seeds.

Explore similar work

ProME: Prototype-Margin Environments with Repair-Aware Selection for Group-Robust Learning
Aug 13, 2026Qianqian Wang, Yunshan Li, Dawei Huang +2Learnable Prototypes Learnability
aShenzhen Key Laboratory of Safety and Security for Next Generation of Industrial Internet, Shenzhen, China · bDepartment of Statistics and Data Science, Southern University of Science and Technology, Shenzhen, China
Discovering Latent Groups for Robust Classification
Jun 22, 2026Ankur Garg, Ulrich Aïvodji, Samira Ebrahimi Kahou +1Decision Trees
Parameter-Efficient Distributionally Robust Adaptation of Tabular Foundation Models under Subpopulation Shift
Oct 1, 2026Seonghwi Kim, Sung Ho Jo, Minwoo ChaeTabular Foundation Models Parameter-Efficient Adaptation
Pohang University of Science and Technology

Aug 13, 2026, cs.LGJ/K move · Enter open · S save

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

Qianqian Wang, Yunshan Li, Dawei Huang, Wenwu Gong+1

aShenzhen Key Laboratory of Safety and Security for Next Generation of Industrial Internet, Shenzhen, China · bDepartment of Statistics and Data Science, Southern University of Science and Technology, Shenzhen, China

Learnable Prototypes Learnability Balance

Open paper AskPDF arXiv

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

Abstract

Explore similar work

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

Discovering Latent Groups for Robust Classification

Parameter-Efficient Distributionally Robust Adaptation of Tabular Foundation Models under Subpopulation Shift