AIA2: Attribute-Agnostic Imbalance Augmentation for Subgroup Robustness

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 (AIA2), a framework for improving model robustness under varying subgroup imbalances without explicit subgroup annotations. AIA2 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 AIA2 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 AIA2 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.
Anthology ID:
2026.findings-emnlp.566
Volume:
Findings of the Association for Computational Linguistics: EMNLP 2026
Month:
October
Year:
2026
Address:
Budapest, Hungary
Editors:
Hung-Yi Lee, Sunipa Dev, Jessy Li, Desmond Elliott
Venue:
Findings
SIG:
SIGDAT
Publisher:
Association for Computational Linguistics
Note:
Pages:
11927–11944
Language:
URL:
https://aclanthology.org/2026.findings-emnlp.566/
DOI:
Bibkey:
Cite (ACL):
Hanshu Rao, Guangzeng Han, and Xiaolei Huang. 2026. AIA2: Attribute-Agnostic Imbalance Augmentation for Subgroup Robustness. In Findings of the Association for Computational Linguistics: EMNLP 2026, pages 11927–11944, Budapest, Hungary. Association for Computational Linguistics.
Cite (Informal):
AIA2: Attribute-Agnostic Imbalance Augmentation for Subgroup Robustness (Rao et al., Findings 2026)
Copy Citation:
PDF:
https://aclanthology.org/2026.findings-emnlp.566.pdf