# AIA2: Attribute-Agnostic Imbalance Augmentation for Subgroup Robustness

arXiv: https://arxiv.org/abs/2608.30297 · HTML: https://mlsanity.com/paper/2608.30297/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 24 of its 47 references that are in the map.
- First posted on arXiv 31 Aug 2026.

## Builds on (most cited first, 24 of 24)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [DeBERTaV3: Improving DeBERTa using ELECTRA-Style Pre-Training with Gradient-Disentangled Embedding Sharing](https://mlsanity.com/paper/2111.09543/) (2021)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [Class-Balanced Loss Based on Effective Number of Samples](https://mlsanity.com/paper/1901.05555/) (2019)
- [EDA: Easy Data Augmentation Techniques for Boosting Performance on Text Classification Tasks](https://mlsanity.com/paper/1901.11196/) (2019)
- [Improving Neural Machine Translation Models with Monolingual Data](https://mlsanity.com/paper/1511.06709/) (2015)
- [Smarter, Better, Faster, Longer: A Modern Bidirectional Encoder for Fast, Memory Efficient, and Long Context Finetuning and Inference](https://mlsanity.com/paper/2412.13663/) (2024)
- [Just Train Twice: Improving Group Robustness without Training Group Information](https://mlsanity.com/paper/2107.09044/) (2021)
- [HateXplain: A Benchmark Dataset for Explainable Hate Speech Detection](https://mlsanity.com/paper/2012.10289/) (2020)
- [Balanced Meta-Softmax for Long-Tailed Visual Recognition](https://mlsanity.com/paper/2007.10740/) (2020)
- [From Louvain to Leiden: guaranteeing well-connected communities](https://mlsanity.com/paper/1810.08473/) (2018)
- [Environment Inference for Invariant Learning](https://mlsanity.com/paper/2010.07249/) (2020)
- [No Subclass Left Behind: Fine-Grained Robustness in Coarse-Grained Classification Problems](https://mlsanity.com/paper/2011.12945/) (2020)
- [Improving Out-of-Distribution Robustness via Selective Augmentation](https://mlsanity.com/paper/2201.00299/) (2022)
- [GPT3Mix: Leveraging Large-scale Language Models for Text Augmentation](https://mlsanity.com/paper/2104.08826/) (2021)
- [Domino: Discovering Systematic Errors with Cross-Modal Embeddings](https://mlsanity.com/paper/2203.14960/) (2022)
- [AugGPT: Leveraging ChatGPT for Text Data Augmentation](https://mlsanity.com/paper/2302.13007/) (2023)
- [Augmenting Data with Mixup for Sentence Classification: An Empirical Study](https://mlsanity.com/paper/1905.08941/) (2019)
- [CrossNER: Evaluating Cross-Domain Named Entity Recognition](https://mlsanity.com/paper/2012.04373/) (2020)
- [HumAID: Human-Annotated Disaster Incidents Data from Twitter with Deep Learning Benchmarks](https://mlsanity.com/paper/2104.03090/) (2021)
- [Examining and Adapting Time for Multilingual Classification via Mixture of Temporal Experts](https://mlsanity.com/paper/2502.08825/) (2025)
- [Diverse Prototypical Ensembles Improve Robustness to Subpopulation Shift](https://mlsanity.com/paper/2505.23027/) (2025)
- [Model-Agnostic Meta Learning for Class Imbalance Adaptation](https://mlsanity.com/paper/2604.18759/) (2026)

## Abstract (arXiv)

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.