Representation Invariance and Allocation: When Subgroup Balance Matters

TL;DR

The paper investigates how subgroup data allocation during training affects subgroup performance when training data are not i.i.d. It introduces the latent separation hypothesis, tying allocation sensitivity to the degree of separation between subgroups in pre-trained latent representations, and formalizes this with a total-variation-based bound. The authors provide theoretical guarantees and extensive empirical evidence across vision and language tasks, showing that representation invariance strongly predicts when allocation will matter, and demonstrate practical data-collection guidance and a TV-based regularisation approach for foundation-model fine-tuning. The work offers a unifying framework that explains when subgroup balancing helps and when it does not, with clear implications for fairness and domain generalisation in real-world deployments.

Abstract

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.

Figures (25)

• 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 balanced accuracy on each subgroup, averaged across 9 fine-tuning runs, is shown alongside estimated linear regression slopes $a$.
• Figure 3: Performance on subgroups under-represented during pre-training (top) and performance on disadvantaged subgroups (bottom) does not necessarily improve with increasing dataset allocation. The y-axis shows the gradient of subgroup loss change with respect to subgroup allocation (negative values indicate performance improvement). No clear correlation is observed in either setting. Each point represents one subgroup, with error bars showing variation 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: AUC 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 AUC difference across 9 fine-tuning runs with error bars indicating standard deviation.

Theorems & Definitions (6)

• Lemma 5.1
• Theorem 5.2: Group accuracy parity
• Lemma E.1
• proof
• Theorem E.2: Group accuracy parity
• proof