AIA$^{2}$: Attribute-Agnostic Imbalance Augmentation for Subgroup Robustness
2026-08-31 • Computation and Language
Computation and Language
AI summaryⓘ
The authors address the problem that some parts of data, like topics or demographics, can be unevenly represented and hurt the model's performance on those less-represented groups. They created a method called Attribute-Agnostic Imbalance Augmentation (AIA²) that finds these hidden imbalances without needing labels for the groups. By using a language model, their method creates more balanced data samples to help the model learn better from difficult and underrepresented subgroups. Tests on various datasets show their approach improves results for the hardest groups compared to other methods.
label imbalancedata attributessubgroup robustnesslatent semantic distributionslarge language modeldata augmentationlearning difficultymodel degradationworst-group performancesubgroup imbalance
Authors
Hanshu Rao, Guangzeng Han, Xiaolei Huang
Abstract
Attributes describing data content and context can induce diverse imbalance patterns that go beyond label imbalance alone. However, existing studies primarily address label imbalance while overlooking data attributes, such as topics and demographics, which can induce meaningful subgroup structure while causing model degradation on underrepresented subgroups. We propose Attribute-Agnostic Imbalance Augmentation (AIA$^{2}$), a framework for improving model robustness under varying subgroup imbalances without explicit subgroup annotations. AIA$^{2}$ automatically discovers varying imbalances via latent semantic distributions, obtains slices with both learning difficulty and subgroup imbalance deficits, and deploys a large language model (LLM) for subgroup-aware imbalance augmentation. We have evaluated AIA$^{2}$ on 5 popular corpora with rich domains and their attribute values, covering social issues and diverse topics. Results show improved performance on the lowest-performing subgroups and consistent gains over competitive baselines. Ablation studies confirm complementary contributions from each component, and additional analyses show that AIA$^{2}$ provides a practical and consistent way to improve worst-group robustness under data subgroup imbalance. Code is available at https://github.com/trust-nlp/AIA2-Subgroup-Robustness.