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

The gist: The degree to which redistributing subgroup data improves subgroup performance is determined by how strongly the set of subgroups are separated in the pre-trained model’s latent representations.

Problem Setting and Motivation

The study investigates how model performance on a particular subgroup scales as its representation in the training data increases, challenging the assumption that balancing subgroup representation optimizes performance. In this work, researchers aim to understand how the allocation of data across subgroups affects subgroup performance at inference time given a fixed training budget. Subgroup allocation is defined as the fraction of the fine-tuning dataset coming from subgroup a(j)k: α(j)k:= 1/K X sum K i=1 1[A(j) i = k], k ∈ [0, 1]>. The central question is how subgroup-specific test performance (for instance the loss l(j)k) depends on these allocations and why this sensitivity differs across attributes and subgroups>.

Empirical Findings on Allocation Sensitivity

The research systematically compares how subgroup allocation affects model performance across various empirical settings. The authors quantify sensitivity in two ways: fitting a linear least-squares regression to subgroup-specific performance as a function of allocation α(j)k, recording the slope a(j)k, and calculating ∆l(j)k = l(j)k (α(j)k = 100%) − l(j)k (α(j)k = 0)>. Across three real-world datasets, they find substantial variability in subgroup performance sensitivity to allocation. Certain subgroups show no benefit from increased allocation and achieve equivalent performance whether the model is fine-tuned only on that subgroup or entirely without it (e.g., age in MIMIC). In contrast, other subgroups, like dataset of origin in HAM10000, are sensitive to their allocation.

Theoretical Framework and Hypothesis Validation

The paper introduces the latent separation hypothesis, stating 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. They formalize this by proving that low latent separation with respect to an attribute A implies low sensitivity to A allocation in the fine-tuning dataset, i.e., fine-tuned model performance on a0 and on a1 does not significantly vary for different allocations α(A). They provide Theorem 5.2, which states that if TV(D') ≤ ε and TV(D'') ≤ ε1, then Accθ'(A = a) − Accθ''(A = a) ≤ 4ε + Accθ' − Accθ'' for all a ∈ [0, 1]>. This theorem links subgroup allocation sensitivity to class-conditional representation separation.

Empirical Evidence and Practical Application

The hypothesis is validated empirically by showing a significant correlation between subgroup separation in pre-trained representations and subgroup sensitivity to allocation during fine-tuning. They find that zero-shot generalisation is directly related to subgroup separation, and the correlation holds across three distance metrics (TV, WD, FD). Furthermore, they demonstrate that enforcing low TV via a differentiable proxy during pre-training reduces subgroup allocation sensitivity. Finally, they apply this practically to foundation models, confirming that balancing by imaging variables significantly increases mean subgroup accuracy when using a 50/50 allocation compared to 100/0 in both models.

Conclusion

The study proposes and proves a novel hypothesis explaining why some subgroup training data allocation does not matter for subgroup performance, and in other cases, it is crucial. By predicting subgroup allocation sensitivity through latent representation analysis, the work provides a new way to inform crucial training data decisions, maximizing fairness, accuracy, and efficiency>. The results support the idea that practitioners should prioritize balancing data with respect to procedure and image view in further fine-tuning data collection efforts. The consistent correlation suggests there may even be a stronger phenomenon at play.

How it works

1. Systematically vary subgroup allocation: The experiment involves varying the allocation α(j)k of different subgroups while keeping the total fine-tuning dataset size K constant.

2. Quantify sensitivity: Sensitivity is quantified by fitting a linear least-squares regression to subgroup-specific performance as a function of allocation α(j)k, recording the slope a(j)k, and calculating ∆l(j)k = l(j)k (α(j)k = 100%) − l(j)k (α(j)k = 0)>.

3. Measure latent separation: Subgroup separation is quantified by extracting penultimate-layer embeddings, projecting them to a lower-dimensional space using PCA, and computing the mean total variation distance (TV) between subgroup distributions conditioned on Y.

4. Test the hypothesis: The core test is whether low latent separation with respect to an attribute A implies low sensitivity to A allocation, which is formally proven by Theorem 5.2. The results show a significant correlation between representation separation and subgroup sensitivity across all three datasets.

5. Operationalize the mechanism: The hypothesis is tested by explicitly enforcing low TV via a differentiable proxy during pre-training, which results in a significant drop in overall performance but reduced sensitivity to subgroup allocation.

1. Data collection strategy informed by latent separation: Practitioners can use quantitative analysis of latent subgroup separation to inform data collection and balancing decisions to prioritize obtaining balanced data with respect to procedure or image view when fine-tuning a foundation model.

2. Model fine-tuning optimization: Fine-tuning should be guided by the degree of separation between subgroups in the pre-trained model’s latent space, as sensitivity to subgroup allocation is significantly correlated to the distance between two subgroups in the pre-trained model’s latent representations.

3. Bias mitigation intervention guidance: The system can identify when data rebalancing will not improve performance by checking if "low latent separation with respect to an attribute A implies low sensitivity to A allocation in the fine-tuning dataset, i.e., fine-tuned model performance on a0 and on a1 does not significantly vary for different allocations α(A)."

4. Foundation model embedding utilization: For foundation models like CheXagent or RAD-DINO-MAIRA-2, selecting a balanced allocation for imaging subgroups can increase subgroup accuracy by over 0.03 when fine-tuning on test sets drawn from different distributions.

5. Representation invariance enforcement: During pre-training, the system can enforce low total variation distance between embedding distributions of subgroups to ensure that group membership does not significantly affect the representation once the label is fixed, leading to reduced sensitivity to allocation changes.

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.