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

Aug 13, 2026cs.LG

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

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.

Qianqian Wang, Yunshan Li, Dawei Huang +2

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

Ask PDF

Jun 22, 2026cs.LG

Discovering Latent Groups for Robust Classification

Machine learning models exploit spurious correlations, achieving high average accuracy but failing disproportionately on underrepresented subgroups. Existing methods address this by adjusting network parameters, guided either by subgroup annotations or inferred pseudo-group labels. Yet at inference, these methods produce only a class prediction, with no insight into a sample's latent subgroup. We propose neural classification trees (NCT), a framework that achieves robustness by encoding subgroup structure in its tree-shaped architecture. By routing each sample to an "easy" or "hard" node of this tree -- based on prediction correctness -- and reusing these routes as pseudo-labels for the next iteration, NCT disentangles conflicting subgroups, without requiring subgroup supervision. We evaluate NCT on five benchmarks spanning binary and multi-class spurious correlations. Our experiments show that the learned tree topology provides strong interpretability by consistently isolating minority subgroups, which provides a transparent mapping between the model architecture and the data's latent group structure, while yielding competitive robustness with state-of-the-art methods.

Ankur Garg, Ulrich Aïvodji, Samira Ebrahimi Kahou +1

Decision Trees

Ask PDF

Oct 1, 2026cs.LG

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

Despite strong mean accuracy, tabular foundation models (TFMs) can perform poorly on underrepresented groups under subpopulation shift, where group proportions change between training and deployment. We propose DR-TFM, a parameter-efficient distributionally robust adaptation framework that requires no true group annotations. DR-TFM adjusts attention to labeled context examples by fine-tuning an existing query scaling network or adding and training one, while keeping all other parameters fixed. We instantiate the framework with two robust objectives using estimated groups or source conditional distributions derived from training data. For TabPFN-3, adaptation updates only 0.016% of the pretrained model's parameters. Across five tabular benchmarks, DR-TFM achieves substantially higher average worst-group accuracy than pretrained TFMs and the compared robust baselines without true group annotations, while maintaining competitive mean group accuracy. DR-TFM also improves average worst-group accuracy on ACS Income and across four additional TFMs.

Seonghwi Kim, Sung Ho Jo, Minwoo Chae

Pohang University of Science and Technology

Tabular Foundation Models Parameter-Efficient Adaptation Distribution Shifts Foundation Model

Ask PDF