When does Subgroup Balance Matter? Predicting Performance Sensitivity from Latent Separability

Anissa Alloula; Charles Jones; Zuzanna Wakefield-Skorniewska; Francesco Quinzan; Bartłomiej Papież

cs.LGDec 10, 2025

When does Subgroup Balance Matter? Predicting Performance Sensitivity from Latent Separability

Authors: Anissa Alloula, Charles Jones, Zuzanna Wakefield-Skorniewska, Francesco Quinzan, Bartłomiej Papież

Organizations: University of Oxford · Imperial College London

Abstract

Unequal representation of demographic groups in training data poses challenges to model generalisation across populations. Standard practice assumes that balancing subgroup representation always optimises performance. However, recent empirical results contradict this assumption: in some cases, imbalanced data distributions actually improve subgroup performance, and in others, subgroup performance is unaffected by the absence of an entire subgroup during training. We propose the latent separation hypothesis, which states that a fine-tuned model's dependence on subgroup representation is determined by the degree of separation between subgroups in the latent space of the pre-trained model. We formalise this hypothesis, provide theoretical analysis, validate it empirically across four vision and language datasets, and present preliminary evidence of a causal role through regularisation experiments. Finally, we show how quantitative analysis of latent subgroup separation can inform practical data collection and balancing decisions when fine-tuning foundation models.

Figures & tables

Figure 1: Model sensitivity to data balance depends on latent separation of subgroups. Left plots show PCA projections of latent representations of MNIST parity classifiers. Right plots show subgroup accuracy as training data allocation changes.

Figure 2: While some subgroups’ accuracy increases with increased representation in training data, others’ performance is independent of their training data representation . The fine-tuned model’s mean balanced accuracy on each subgroup across 9 runs and regression slopes $a$ are shown.

Figure 3: Performance on subgroups under-represented during pre-training (left), disadvantaged (middle), and with class imbalance (right) does not necessarily improve with increasing dataset allocation . No clear correlation to performance change (a negative gradient of subgroup loss change indicates performance improvement) is observed in any of the 3 settings. Each point represents one subgroup, with error bars showing standard deviation across 9 runs.

Figure 4: Sensitivity to subgroup allocation is highly correlated with separation in the pre-trained model’s representation space (as measured by total variation distance, TV) across the three datasets. Each dot represents mean TV and loss slope for one subgroup, averaged across 9 fine-tuning runs, with bars corresponding to standard deviations, and Pearson correlation also shown.

Figure 5: Subgroup generalisation degrades as latent subgroup separation increases. Mean and standard deviation of balanced accuracy difference when a subgroup is excluded from training are shown across 3 training-from-scratch runs.

Figure 6: TV regularisation reduces subgroup allocation sensitivity ; subgroup accuracy is constant.

Figure 7: In foundation model fine-tuning , balanced allocations increase subgroup accuracy by 0.02 for imaging groups, but yield little improvement for demographic groups, as predicted by their reduced TV.

Appendix figures & tables28 assets

Supplementary material from the paper’s appendix.

Appendix

Training strategy	MNIST	MIMIC-CXR	HAM10000	Civil_comments
$Y$	Even/odd digit	Pleural effusion	Malignant/benign lesion	Comment toxicity
Backbone	2-layer CNN	DenseNet121 Huang et al. (2016)	ViTB16 Dosovitskiy et al. (2021)	BERTClassifier (uncased) Devlin et al. (2018)
Initialisation	None	ImageNet Deng et al. (2009)	ImageNet Deng et al. (2009)	Bookcorpus, Wikipedia (English)
$\mathbf{N_{pretrain}}$	5,000	33,237	2,000	17,920
$\mathbf{N_{finetune}}$	24,000	22,353	3,184	24,000
$\mathbf{N_{test}}$	4,000	26,590	1,000	12,243

Appendix

Table 1: Implementation details for all models.

Dataset	Attribute	Group 0 / 1	$\mathbf{P(A{=}1)}$	$\mathbf{P(Y\|A{=}0)/P(Y\|A{=}1)}$
MNIST	Colour	Red / Green	0.50	0.50 / 0.50
	Digit value	Below 5 / Over	0.50	0.50 / 0.50
	Random	/	0.50	0.50 / 0.50
MIMIC-CXR	View Position	Lateral / Frontal	0.65	0.16 / 0.26
	Patient Orientation	Recumbent / Erect	0.77	0.35 / 0.20
	Procedure	Portable / Fixed scanner	0.46	0.35 / 0.11

Appendix

Table 2: Subgroups used for each dataset and their base population prevalence and class prevalence.

Dataset	Attribute	Loss slope	Balanced Acc slope	AUC slope
MIMIC	View Position	-0.027 (0.001)	0.034 (0.003)	0.012 (0.001)
MIMIC	Patient Orientation	-0.009 (0.001)	0.013 (0.002)	0.003 (0.000)
MIMIC	Procedure	-0.026 (0.001)	0.026 (0.002)	0.016 (0.001)
MIMIC	Support_Devices	-0.009 (0.001)	0.010 (0.001)	0.006 (0.000)
MIMIC	Gender	-0.003 (0.001)	0.004 (0.002)	0.002 (0.000)
MIMIC	Insurance	-0.002 (0.001)	0.003 (0.002)	0.001 (0.000)

Appendix

Table 3: Sensitivity of subgroup performance to fine-tuning allocation across datasets and attributes in full model fine-tuning. Reported values are esimated regression slopes (mean

\pm

standard deviation across nine fine-tuning runs) for subgroup loss, balanced accuracy, and AUC as allocation varies from 0% to 100%. Larger absolute slopes indicate higher sensitivity.

Dataset	Attribute	Loss slope	Balanced Acc slope	AUC slope
MIMIC	View Position	-0.007 (0.001)	0.001 (0.002)	-0.000 (0.000)
MIMIC	Patient Orientation	-0.008 (0.001)	-0.002 (0.001)	0.000 (0.000)
MIMIC	Procedure	-0.011 (0.001)	0.014 (0.002)	0.003 (0.000)
MIMIC	Support_Devices	-0.005 (0.001)	0.000 (0.001)	0.000 (0.001)
MIMIC	Gender	-0.000 (0.001)	0.001 (0.002)	0.000 (0.000)
MIMIC	Insurance	-0.000 (0.001)	0.000 (0.002)	-0.000 (0.000)

Appendix

Table 4: Sensitivity of subgroup performance to fine-tuning allocation across datasets and attributes in last-layer fine-tuning. Reported values are esimated regression slopes (mean

\pm

standard deviation across nine fine-tuning runs) for subgroup loss, balanced accuracy, and AUC as allocation varies from 0% to 100%. Larger absolute slopes indicate higher sensitivity.

Attribute	Group	$\sigma_{g}$	$p_{g}$	$\tau_{g}$	$q_{g}$	$\delta_{g}$
Gender	0	0.000 $\pm$ 0.000	1.298 $\pm$ 0.000	3136.744 $\pm$ 13745.833	2.000 $\pm$ 0.964	0.302 $\pm$ 0.020
Gender	1	278.269 $\pm$ 647.492	2.000 $\pm$ 0.650	3002.106 $\pm$ 13853.293	2.000 $\pm$ 1.015	0.348 $\pm$ 0.019
ViewPosition	0	4.843 $\pm$ 6.518	0.850 $\pm$ 0.387	1.523 $\pm$ 4.728	0.225 $\pm$ 1.188	0.000 $\pm$ 1.388
ViewPosition	1	1.035 $\pm$ 0.806	0.260 $\pm$ 0.336	2.175 $\pm$ 8.930	0.276 $\pm$ 1.359	0.000 $\pm$ 1.609

Appendix

Table 5: Estimated power-law scaling model parameter fits and standard deviations.

Dataset	No. of groups	Spearman $\rho$	$p$ -value
MIMIC	10	0.77	0.009
HAM10000	5	0.90	0.037
CivilComments	6	0.83	0.042

Appendix

Table 6: Correlation between full-range and extreme-allocation sensitivity estimates. Sensitivity is measured as the mean slope of the loss across the 0-10% and 90-100% allocation range (the extremes) and compared to the loss slope across the full allocation range.

Figure 8: Sensitivity to subgroup allocation at the thresholds is highly correlated with separation in the pre-trained model’s representation space (as measured by total variation distance, TV) across the three datasets. Each dot represents mean TV and loss slope (between 0 and 10% and 90 and 100% subgroup allocation) for one subgroup, averaged across 9 fine-tuning runs, with bars corresponding to standard deviations, and Pearson correlation also shown.

Figure 9: Performance on subgroups under-represented during pre-training (top), disadvantaged (middle), and with class imbalance (bottom) does not necessarily improve with increasing dataset allocation. No clear correlation to performance change (a negative gradient of subgroup loss change indicates performance improvement) is observed in any of the three settings. Each dot represents one subgroup with bars indicating variation across 9 fine-tuning runs.

Dataset	Hypothesis	Correlation	$p$ -value	Significant ( $p<0.05$ )	$N$
MIMIC	Subgroup prevalence	0.0586	0.7955	False	22
MIMIC	Subgroup accuracy	0.5197	0.01318	True	22
MIMIC	Class imbalance	-0.1906	0.3954	False	22
HAM10000	Subgroup prevalence	-0.1027	0.7777	False	10
HAM10000	Subgroup accuracy	0.0616	0.8658	False	10
HAM10000	Class imbalance	0.0261	0.9430	False	10

Appendix

Table 7: Correlation results across datasets and hypotheses.

Figure 10: Overall and group-wise balanced accuracy across all fine-tuning runs for MNIST, MIMIC, HAM10000, and Civil_comments.

Figure 11: Value of upper bound for MNIST, MIMIC, HAM10000, and Civil_comments subgroups, with the trivial upper bound (accuracy difference of 1) marked in red.

Figure 12: Total variation (TV) distances across subgroups in the pre-trained model’s penultimate-layer representation space. We report mean and standard deviation across three random seeds.

Figure 13: Total variation distance (TV), wasserstein distance (WD), and Fréchet distance (FD) across subgroups appear correlated in MIMIC. Each dot represents a subgroup with error bars representing standard deviation across three pre-training runs.

Figure 14: Sensitivity to subgroup allocation is correlated with separation in the pre-trained model’s representation space (as measured by wasserstein distance, and Fréchet distance) across the three datasets. Each dot represents mean distance and loss slope for one subgroup, averaged across 9 fine-tuning runs, with bars corresponding to standard deviations, and Pearson correlation coefficients also shown.

Figure 15: Accuracy difference when a model has not been trained on a subgroup increases with the separation of the latent representations of the subgroups. Results are shown as mean balanced accuracy difference across 9 fine-tuning runs with error bars indicating standard deviation.

Figure 16: Sensitivity to subgroup allocation (slope of AUC) is correlated with subgroup TV distance in the pre-trained model’s representation space across the three datasets. Each dot represents mean distance and AUC slope for one subgroup, averaged across 9 fine-tuning runs, with bars corresponding to standard deviations, and Pearson correlation coefficients also shown.

Figure 17: Examples showing latent representation separations are stable across subgroup allocations under full fine-tuning and close to separation of the initial pre-trained model.

Figure 18: Correlation between loss slope and distance metrics remains strong under full fine-tuning. Each dot represents mean total variation distance and loss slope for one subgroup, averaged across 9 fine-tuning runs, with bars corresponding to standard deviations, and Pearson correlation coefficients also shown.

Figure 19: Representation distances continue to predict accuracy gaps at extreme subgroup allocations under full fine-tuning at 100% allocation vs 0%. Bars represent mean and standard deviation of the balanced accuracy across 9 fine-tuning runs.

Figure 20: Correlation between loss slope and distance metrics remains strong when training from scratch. Each dot represents mean total variation distance and loss slope for one subgroup, averaged across 3 training runs, with bars corresponding to standard deviations, and Pearson correlation coefficients also shown.

Figure 21: Total variation distances of CheXagent (left) and MAIRA-2 (right) embeddings of MIMIC images across subgroups.

Figure 22: In RAD-DINO-MAIRA-2 foundation model fine-tuning, selecting a balanced allocation for imaging subgroups increases subgroup accuracy by over 0.03, while it has less importance for demographic subgroups, as predicted by their reduced total variation distance. We report mean accuracy and standard deviations for 3 fine-tuning runs, with red arrows for potential gains from balanced allocations.

Figure 23: Correlation between latent representation separation and subgroup allocation sensitivity in CheXagent (vision-language model) fine-tuning. Results confirm that our hypothesis generalises beyond task-specific models. Each dot represents mean distance and loss slope for one subgroup, averaged across 9 fine-tuning runs, with bars corresponding to standard deviations, and Pearson correlation coefficients also shown.

Figure 24: Correlation between latent representation separation and subgroup allocation sensitivity in MAIRA-2 (vision-language model) fine-tuning. Results confirm that our hypothesis generalises beyond task-specific models. Each dot represents mean distance and loss slope for one subgroup, averaged across 9 fine-tuning runs, with bars corresponding to standard deviations, and Pearson correlation coefficients also shown.

Figure 25: In CheXagent foundation model parameter-efficient fine-tuning, selecting a balanced allocation for imaging subgroups increases subgroup accuracy by about 0.03, while it has less importance for demographic subgroups, as predicted by their reduced total variation distance. We report mean accuracy and standard deviations for 3 fine-tuning runs, with red arrows for potential gains from balanced allocations.

Figure 26: Correlation between latent representation separation and subgroup allocation sensitivity in CheXagent (vision-language model) parameter-efficient fine-tuning. Results confirm that our hypothesis generalises beyond task-specific models. Each dot represents mean distance and loss slope for one subgroup, averaged across 9 fine-tuning runs, with bars corresponding to standard deviations, and Pearson correlation coefficients also shown.

Figure 27: Correlation between representation separation and subgroup allocation sensitivity is strong across fine-tuning dataset sizes in MIMIC. Each dot represents mean distance and loss slope for one subgroup, averaged across 9 fine-tuning runs, with bars corresponding to standard deviations, and Pearson correlation coefficients also shown.

Method	TV (mean $\pm$ std)	Overall acc. (mean $\pm$ std)	Subgroup acc. change (mean $\pm$ std)
Baseline	$0.995\pm 0.006$	$0.794\pm 0.004$	$0.010\pm 0.002$
Regularised	$0.87\pm 0.012$	$0.737\pm 0.006$	$-0.000\pm 0.002$

Appendix

Table 8: Full results from regularisation experiments, including mean total variation distance (TV) between MIMIC subgroups, overall accuracy, and subgroup accuracy change across allocations.

Subgroup Robustness

See all 10 papers

AIA $^{2}$ : Attribute-Agnostic Imbalance Augmentation for Subgroup Robustness
Aug 31, 2026Hanshu Rao, Guangzeng Han, Xiaolei HuangClass-Imbalanced Learning Subgroup Robustness
University of Memphis · University of Missouri
Subgroup performance analysis of adaptation strategies for chest X-ray foundation models
Aug 19, 2026Dhruv Gupta, Emma A. M. Stanley, Fabio De Sousa Ribeiro +2Chest X-Ray Classification Medical Imaging Foundation Models
Department of Computing, Imperial College London, UK · Causality in Healthcare AI Hub, UK · Department of Radiology, Royal Brompton Hospital, UK +1
When Data Imbalance Helps: Robust Generalization Through Shortcut Saturation
Jul 11, 2026Cheng-Ting Chou, Duc Binh HoangClass-Imbalanced Learning Transformer
University of California, Los Angeles

Aug 31, 2026J/K move · Enter open · S save

AIA $^{2}$ : Attribute-Agnostic Imbalance Augmentation for Subgroup Robustness

Hanshu Rao, Guangzeng Han, Xiaolei Huang

University of Memphis · University of Missouri

Class-Imbalanced Learning Subgroup Robustness Synthetic Data Augmentation

Open paper AskPDF arXiv

Attributes describing data content and context can induce diverse imbalance patterns that go beyond label imbalance alone. However, existing studies primarily address label imbalance while overlooking data attributes, such as topics and demographics, which can induce meaningful subgroup structure while causing model degradation on underrepresented subgroups. We propose Attribute-Agnostic Imbalance Augmentation (AIA

^{2}

), a framework for improving model robustness under varying subgroup imbalances without explicit subgroup annotations. AIA

^{2}

automatically discovers varying imbalances via latent semantic distributions, obtains slices with both learning difficulty and subgroup imbalance deficits, and deploys a large language model (LLM) for subgroup-aware imbalance augmentation. We have evaluated AIA

^{2}

on 5 popular corpora with rich domains and their attribute values, covering social issues and diverse topics. Results show improved performance on the lowest-performing subgroups and consistent gains over competitive baselines. Ablation studies confirm complementary contributions from each component, and additional analyses show that AIA

^{2}

provides a practical and consistent way to improve worst-group robustness under data subgroup imbalance. Code is available at https://github.com/trust-nlp/AIA2-Subgroup-Robustness.

Abstract

Subgroup Robustness

Explore similar work

AIA2^{2}2: Attribute-Agnostic Imbalance Augmentation for Subgroup Robustness

Subgroup performance analysis of adaptation strategies for chest X-ray foundation models

When Data Imbalance Helps: Robust Generalization Through Shortcut Saturation

AIA $^{2}$ : Attribute-Agnostic Imbalance Augmentation for Subgroup Robustness