Guiding Data Allocation for Robust Subpopulation Generalization

Zhaoying Pan; Yipei Wang; Shenyu Lu; Xiaoqian Wang

cs.LGNov 9, 2026

Guiding Data Allocation for Robust Subpopulation Generalization

Authors: Zhaoying Pan, Yipei Wang, Shenyu Lu, Xiaoqian Wang

Organizations: Elmore Family School of Electrical and Computer Engineering, Purdue University

Abstract

Machine learning models often achieve high average accuracy in the training data while performing poorly on test data due to subpopulation shift. A common data-centric solution is to enforce balanced subgroup proportions, but full balance can be costly, infeasible, and not always necessary. In this work, we study whether balanced training data is the unique optimal configuration for robust subpopulation generalization. We extend a Gaussian feature-model framework to theoretically characterize the performance landscape across general training subgroup distributions. We show that multiple configurations, including substantially imbalanced ones, can perform comparably to the balanced one. We further derive gradient-based directions for adjusting subgroup proportions to improve adjusted accuracy and worst-group accuracy more efficiently. These findings suggest that strategically allocating data offers a more flexible and principled path to robust performance than simple balancing. The code is available at https://github.com/zhaoyingpan/subpop-data-allocation.

Figures & tables

Panel of Figure 1 — Figure 1: (a) Test accuracy varies over diverse training data configurations and many match or exceed the performance of the balanced configuration. (b) Data configurations visualized by projecting group sizes onto a 3-simplex, including the configurations for the balanced (center), the highly skewed (vertex), and a random configuration (unique coordinate). (c) Visualization of directions estimated according to the gradients (Grad) and balancing (Bal) towards the center.

Panel of Figure 4 — Figure 4: Existence of divergent directions across datasets. Empirical comparison of Grad/Bal directions with accuracy gaps obtained from $7^{3}=343$ configurations of training data. Lighter colors are observed near edges connecting $r_{-1,-1}$ to $r_{-1,1}$ , and $r_{1,-1}$ to $r_{1,1}$ , indicating a greater disparity between Grad/Bal directions and aligning with the theoretical accuracy gap trends in Figure 3 .

Panel of Figure 5 — Figure 5: Efficiency of data allocation along directions across datasets and initial configurations. The efficiency of directions is validated by the comparison of the performance change as the data is allocated according to various directions. Each row corresponds to one allocation experiment with an initial configuration, and each column corresponds to a dataset. The legends are shared across plots.

Panel of Figure 6 — Figure 6: Optimality of test performance across multiple configurations. Among $7^{3}=343$ training data configurations, those achieving test performance above the 90th percentile are visualized within the 3-simplex. The results confirm that multiple configurations can achieve test performance comparable to or exceeding that of the balanced configuration, consistent with Figure 3 .

Method	Adjusted Acc.	WGA
Original	0.766	0.376
Over-sampling	0.768	0.366
Over-sampling + Grad	0.821	0.541
Under-sampling	0.777	0.837
Under-sampling + Grad	0.901	0.848

Table 1: Grad-guided resampling on Example 2 of Waterbirds-syn . Vanilla over-/under-sampling targets the balanced configuration, while Grad-guided resampling targets the alternative optimal configurations.

		Waterbirds		CelebA
Method	Metric	Original	Grad	Original	Grad
DFR	WGA	89.30 $\pm$ 0.55	89.10 $\pm$ 0.75	87.49 $\pm$ 1.27	87.36 $\pm$ 0.13
DFR	AA	91.37 $\pm$ 0.45	91.60 $\pm$ 0.19	90.87 $\pm$ 0.82	90.39 $\pm$ 0.33
SELF	WGA	88.28 $\pm$ 1.37	88.53 $\pm$ 0.76	82.59 $\pm$ 2.31	82.18 $\pm$ 3.81
SELF	AA	91.50 $\pm$ 0.26	91.74 $\pm$ 0.26	90.52 $\pm$ 0.57	90.74 $\pm$ 0.72
GSR	WGA	91.90 $\pm$ 1.42	91.96 $\pm$ 0.95	85.74 $\pm$ 2.57	86.67 $\pm$ 0.96

Table 2: Comparison of worst-group and adjusted accuracy (%) under different stage-two data allocation strategies. Results are reported for last-layer retraining using the original stage-two allocation and the Grad allocation.

Appendix figures & tables25 assets

Supplementary material from the paper’s appendix.

Appendix

	Waterbirds-syn	UrbanCars-syn	CelebA-subset	CivilComments-subset
$m_{1}$	2.142	0.567	1.903	0.977
$m_{2}$	8.0	0.789	2.0	2.0

Appendix

Table 3: Estimated Separability for Datasets

	$\alpha_{+}$	$\alpha_{-}$	$\beta$
Example 1	0.848	0.695	0.059
Example 2	0.976	0.184	0.128
Example 3	0.397	0.883	0.862

Appendix

Table 4: Initial Configurations for Data Allocation Experiments in Section 4.3

Figure 9: Comparison of the Grad and Bal directions for initial configurations in the 3-simplex.

Panel of Figure 10 — Figure 10: Optimality of theoretical test performance.

Panel of Figure 11 — Figure 11: Optimality of test performance (95th percentile) across multiple configurations.

Panel of Figure 12 — Figure 12: Optimality of test performance (99th percentile) across multiple configurations.

Panel of Figure 13 — Figure 13: Efficiency of WGA of data allocation along directions across datasets and initial configurations.

Panel of Figure 14 — Figure 14: Optimality of test WGA across multiple configurations.

	Waterbirds-syn	UrbanCars-syn	CelebA-subset	CivilComments-subset
Example 1	84.15 / 79.68	57.87 / 50.93	80.19 / 69.26	59.70 / 54.46
Example 2	73.24 / 65.79	40.80 / 29.87	63.98 / 51.79	32.30 / 30.12
Example 3	90.47 / 88.41	49.33 / 45.87	76.94 / 75.54	61.61 / 51.42

Appendix

Table 5: WGA (%) after 15 allocation steps. Each cell reports Grad / Bal (mean over three seeds).

	Waterbirds-syn	UrbanCars-syn	CelebA-subset	CivilComments-subset
Example 1	19.95 / 19.41	12.30 / 10.95	18.35 / 16.76	13.63 / 12.32
Example 2	17.00 / 14.88	8.21 / 6.06	13.53 / 11.07	7.82 / 6.95
Example 3	21.84 / 21.44	11.12 / 10.21	18.02 / 17.74	14.86 / 12.96

Appendix

Table 6: WGA-AUC over the allocation trajectory. Each cell reports Grad / Bal (mean over three seeds).

Panel of Figure 15 — Figure 15: Relationship between Adjusted Accuracy, WGA, and Accuracy Gap.

Direction	AA@15	WGA@15	AA-AUC	WGA-AUC
AA-derived	$92.45{\pm}0.29$	$83.93{\pm}1.03$	$22.81{\pm}0.09$	$20.05{\pm}0.24$
WGA-derived	$92.37{\pm}0.60$	$84.29{\pm}1.55$	$22.77{\pm}0.12$	$19.99{\pm}0.36$

Appendix

Table 7: AA-derived versus WGA-derived allocation on Waterbirds-syn Example 1. Results are reported as mean

\pm

standard deviation over three seeds.

Method	AA	WGA	AA-AUC	WGA-AUC
Scaling-law allocation	91.78	79.38	22.69	19.40
Grad (ours)	92.50	84.15	22.78	19.95

Appendix

Table 8: Comparison with the adapted scaling-law data-allocation baseline Rolf et al. (2021) on Waterbirds-syn. AA and WGA are measured after 15 allocation steps.

Method	Group 1	Group 2	Group 3	Group 4
Original	1124	28	6404	1444
Over-sampling	6404 (+5280)	6404 (+6376)	6404 (+0)	6404 (+4960)
Over-sampling (Grad)	5799 (+4675)	857 (+829)	6404 (+0)	811 (-633)
Under-sampling	28 (-1096)	28 (+0)	28 (-6376)	28 (-1416)
Under-sampling (Grad)	190 (-934)	28 (+0)	209 (-6195)	27 (-1417)

Appendix

Table 9: Group sizes for the Waterbirds-syn data variants used in the resampling experiment. Values in parentheses indicate the change relative to the original group size.

	Waterbirds-syn	UrbanCars-syn	CelebA-subset	CivilComments-subset
Example 1	.055 / .062	.066 / .028	.003 / .003	.056 / .060
Example 2	.012 / .0001	.012 / .005	.001 / .002	.246 / .010
Example 3	.037 / .017	.081 / .024	.007 / .030	.002 / .011

Appendix

Table 10:

p

-values for Grad versus Bal in AA at step 15 / AA-AUC over the allocation trajectory (three seeds).

Figure 16: Integration of Group DRO and Regularization with Our Method of Allocating Data.

Figure 17: Grad Robustness to Separability Estimation Error (Overall).

Panel of Figure 18 — Figure 18: Grad Robustness to Separability Estimation Error. We vary the estimated separability by 0.9, 1, and 1.1 to synthesize the estimation error

	Accuracy Range	Accuracy at Balanced Setup
ResNet50	0.959 - 0.966	0.953
ViT	0.958 - 0.973	0.961

Appendix

Table 11: Optimality experiments on Waterbirds-syn with top 10% configurations on ResNet50.

Panel of Figure 19 — Figure 19: Data Allocation on UrbanCars-syn with Three Attributes.

Figure 20: Efficiency of data allocation along directions across initial configurations on Waterbirds-syn with ViT

Figure 21: Efficiency of data allocation along directions across initial configurations on CelebA-subset with ViT

Figure 22: Efficiency of data allocation along directions across initial configurations on CivilComments-subset with DeBERTa

Dataset	Accuracy Range
Waterbirds-syn	0.954 $\pm$ 0.003 – 0.965 $\pm$ 0.002
Urbancars-syn	0.826 $\pm$ 0.007 – 0.844 $\pm$ 0.007
CelebA-subset	0.893 $\pm$ 0.003 – 0.902 $\pm$ 0.002
CivilComments-subset	0.796 $\pm$ 0.003 – 0.807 $\pm$ 0.003

Appendix

Table 12: Optimality experiments with 3 random seeds on configurations with top performance.

Subgroup Robustness

See all 12 papers

When does Subgroup Balance Matter? Predicting Performance Sensitivity from Latent Separability
Dec 10, 2025Anissa Alloula, Charles Jones, Zuzanna Wakefield-Skorniewska +2Subgroup Robustness Dataset Bias
University of Oxford · Imperial College London
Active Learning of Robust Models
Nov 6, 2026Mizu Nishikawa-Toomey, Andreas Kirsch, Laurent Charlin +1Active Learning Distribution Shift Robustness
Mila – Quebec AI Institute · Université de Montréal · HEC Montréal
Rank-Constrained Adaptation for Reliable Real-World Performance
Feb 6, 2026Abinitha Gourabathina, Hyewon Jeong, Teya Bergamaschi +2Subgroup Robustness Distribution Shift Robustness
Department of Electrical Engineering & Computer Science, Massachusetts Institute of Technology, Cambridge, MA, USA · Harvard School of Medicine, Harvard, Boston, MA, USA.

Dec 10, 2025J/K move · Enter open · S save

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

Anissa Alloula, Charles Jones, Zuzanna Wakefield-Skorniewska, Francesco Quinzan+1

University of Oxford · Imperial College London

Subgroup Robustness Dataset Bias Fine-Tuning Class-Imbalanced Learning

Open paper AskPDF arXiv

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.

Abstract

Subgroup Robustness

Explore similar work

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

Active Learning of Robust Models

Rank-Constrained Adaptation for Reliable Real-World Performance