Strategic data allocation improves robustness beyond balanced groups
Guiding Data Allocation for Robust Subpopulation Generalization
Machine LearningComputer Vision and Pattern Recognition
Summary
Machine learning models often do well on average with training data but struggle when tested on groups that differ from their training examples. A usual fix is to balance data across all subgroups, but this can be expensive and not always needed. The authors found that unevenly sized groups can still perform just as well, and they developed ways to adjust data proportions to improve accuracy more efficiently. Their approach offers more flexibility than just balancing all subgroups equally.
What this means in practice
- •For machine learning engineers: Adjust training data proportions strategically to improve model robustness against subgroup shifts rather than enforcing full balance.
- •For data labeling teams: Allocate labeling efforts more efficiently by prioritizing subgroup sizes based on gradient guidance instead of uniformly balancing all subgroups.
Authors
Zhaoying Pan, Yipei Wang, Shenyu Lu, Xiaoqian Wang
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.