Representation Invariance and Allocation: When Subgroup Balance Matters
This addresses the problem of optimizing model generalization across populations for AI practitioners, offering insights into when subgroup balance matters, though it is incremental in refining existing practices.
The paper investigates how varying the representation of demographic subgroups in training data affects model performance, finding that imbalanced data can sometimes improve subgroup performance and that performance may not depend on the presence of entire subgroups. It proposes the latent separation hypothesis, validated empirically, to explain these effects and applies it to guide data balancing in foundation model fine-tuning.
Unequal representation of demographic groups in training data poses challenges to model generalisation across populations. Standard practice assumes that balancing subgroup representation optimises performance. However, recent empirical results contradict this assumption: in some cases, imbalanced data distributions actually improve subgroup performance, while in others, subgroup performance remains unaffected by the absence of an entire subgroup during training. We conduct a systematic study of subgroup allocation across four vision and language models, varying training data composition to characterise the sensitivity of subgroup performance to data balance. We propose the latent separation hypothesis, which states that a partially 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, and validate it empirically. Finally, we present a practical application to foundation model fine-tuning, demonstrating that quantitative analysis of latent subgroup separation can inform data collection and balancing decisions.