PaperlayerView exact revision on arXiv

Source-linked AI summary

AIA$^{2}$: Attribute-Agnostic Imbalance Augmentation for Subgroup Robustness

Hanshu Rao, Guangzeng Han, Xiaolei Huang

arXiv:2608.30297v1cs.CL

TL;DR

Existing imbalance methods often overlook attribute-induced subgroup shifts, leaving robustness concerns for underrepresented groups. AIA2 discovers latent slices, analyzes local label deficits, and uses targeted LLM augmentation without explicit subgroup annotations. Across five datasets, it improves overall and worst-group performance over competitive baselines, while its computational cost and generalizability remain bounded concerns.

  • Problem

    Existing studies primarily address global label imbalance while overlooking attribute-induced subgroup imbalance, which can produce unstable performance across underrepresented subgroups.

  • Method

    AIA2 iteratively discovers slices in a joint semantic-predictive space, identifies local label deficits, and applies constraint-guided LLM augmentation to targeted hard examples.

  • Results

    Across five text classification and named entity recognition datasets, AIA2 improves overall and worst-group performance over competitive baselines.

  • Takeaways & Limitations

    AIA2 provides a practical way to improve worst-group robustness under subgroup imbalance without explicit attribute annotations.

  • Takeaways & Limitations

    AIA2 adds computational cost, depends on reliable semantic and prediction representations, and requires further validation on other NLP tasks.

Abstract

from arXiv · show

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.

1 Introduction

Data attributes create subgroup imbalance patterns beyond global label skew, while existing approaches often overlook these shifts and can leave minority-group performance unstable. AIA2 addresses this by discovering latent slices, analyzing distribution gaps, and using guided augmentation to improve robustness without explicit subgroup metadata.

  • Motivation: Data attributes such as topics, demographics, and sources create subgroups with distinct statistical properties and diverse imbalance patterns.
  • Motivation: Models optimized for overall label distributions can still show substantial performance disparities across attribute-defined subgroups.
  • Motivation: Event-source subgroups in HumAID exhibit label imbalance with substantially varying imbalance patterns across sources.
  • AIA2 framework: Constraint-guided LLM augmentation generates and filters diverse, consistent candidates before merging them into training for the next iteration.
  • AIA2 framework: AIA2 discovers latent slices using joint semantic and predictive representations, then identifies distributional deficits and high-loss seeds.
  • AIA2 framework: The framework targets subgroup imbalance without predefined attribute values or explicit subgroup labels during training or model selection.

2 Method

AIA2 iteratively discovers latent slices using semantic and predictive information, identifies local label deficits, and augments difficult examples without predefined subgroup labels. It combines constrained LLM generation with filtering to repair distribution defects as model behavior evolves.

  • 2.1 Latent Slice Discovery: AIA2 discovers latent slices by combining semantic embeddings with predictive behavior, capturing coherent regions that single-distribution clustering can miss.A k-nearest neighbor graph and Leiden community detection identify variable-sized communities at multiple scales.
  • 2.1 Latent Slice Discovery: The joint representation concatenates L2-normalized semantic and predictive vectors, with α controlling the relative importance of predictive information.For NER, predictive vectors encode boundary losses, entropy, and negative margins across entity types.
  • 2.2 Slice-Level Distribution Gap Analysis: Within each slice, AIA2 compares local and global label distributions and selects slices whose L1 distribution gap exceeds threshold τ.The deficit vector identifies labels under-represented relative to the global distribution.
  • 2.2 Slice-Level Distribution Gap Analysis: Augmentation seeds are selected from deficit labels by prioritizing the highest-loss examples, with seed counts proportional to deficit mass.The slice budget uses sublinear scaling, an overall augmentation rate ρ, and a per-slice cap nmax.
  • 2.3 Constraint-Guided Augmentation and Filtering: An LLM generates label-preserving, diverse candidates using structural transformations, type-preserving entity substitutions, and masked reference examples.NER generation preserves placeholders and restores original entities with corresponding BIO tags.
  • 2.3 Constraint-Guided Augmentation and Filtering: Candidates pass hard format checks and are scored for prediction consistency and semantic diversity before retaining typically 2–3 augmentations per seed.AIA2’s unified training procedure repeatedly discovers slices, selects high-loss deficit-label seeds, and generates augmentations.

3 Experiment

AIA2 is evaluated on five publicly available corpora spanning text classification and named-entity recognition, against baselines covering reweighting, group optimization, and LLM-based augmentation. The experiments use fixed task and augmentation models, iterative training, multiple metrics, and five random seeds.

  • Datasets and tasks: Five publicly available corpora are used to evaluate AIA2 across text classification and named-entity recognition.The datasets include globally imbalanced classes and heterogeneous subpopulations whose local label distributions may differ from global patterns.
  • Baselines: AIA2 is compared with sample-level reweighting, group-level optimization, and LLM-based augmentation baselines.The compared methods include Focal Loss, JTT, Group DRO, DPE, GEORGE, AugGPT, and CB-LLM.
  • Baselines: AIA2 performs targeted augmentation by inferring latent subgroup-label structure, unlike baselines that use global class distributions or random paraphrasing.The framework targets augmentation according to inferred latent subgroup-label structure.
  • Experimental settings: DeBERTa-v3-base is the default task model, while BGE-large-en-v1.5 discovers slices and Qwen3-32B performs sample augmentation.Qwen3-32B is used with temperature 0.85; training runs for eight iterative rounds and selects the best checkpoint by validation performance.
  • Evaluation: Performance is reported with Micro-F1 and Macro-F1, plus worst-group accuracy for classification and worst-group F1 for named-entity recognition.For NER, scores are computed at entity level under BIO, excluding the O tag.
  • Evaluation: All results are averaged over five independent runs with different random seeds, with standard deviations reported in an appendix.Additional appendices cover implementation details, hyperparameter sensitivity, and computational overhead.

4 Results

AIA2 improves overall and subgroup-level robustness across five datasets, with gains linked to repairing local label deficits. Ablations and sensitivity analyses support the contributions and generality of latent slice discovery, gap-aware selection, and LLM augmentation.

  • Overall performance: AIA2 improves Macro-F1, Micro-F1, and worst-group metrics by up to 1.9, 2.2, and 1.5 points, respectively, across datasets.On WWW2015, it reaches 57.6 Macro-F1 and improves WGA over the best baseline by 1.5 points.
  • Subgroup-level analysis: AIA2 delivers broad subgroup gains across all five datasets, forming the outer radar envelope on most subgroup spokes without trading off specific subgroups.Improvements are especially evident on challenging minority, event-specific, and domain-specific subgroups.
  • Slice–attribute alignment: Discovered-slice alignment varies across datasets, with AMI values of 0.637 on HumAID, 0.615 on WWW2015, 0.559 on CrossNER, and 0.268 on RE3D.Despite weaker alignment on RE3D, AIA2 improves all three evaluation metrics there, indicating that useful slices need not reproduce metadata groups.
  • Deficit repair and subgroup gains: On WWW2015, deficit reduction correlates with accuracy gain across 39 subgroup-label subsets, with Spearman correlation 0.769 and OLS R2 = 0.487.Sixteen subsets, or 41.0%, fall in the region with both reduced deficit and improved accuracy.
  • Module contributions: Removing latent slice discovery, gap-aware seed selection, or LLM augmentation degrades performance, including WWW2015 Macro-F1 drops from 57.6 to 54.3, 53.8, or lower.The largest reported gap-aware effect is RE3D WG-F1 dropping from 60.5 to 53.3.
  • Robustness to base models: AIA2 remains strongest or comparable across ModernBERT and multiple LLM generators, though WWW2015 varies by 3.6 Macro-F1 points between Qwen3-32B and Qwen3-14B.On CrossNER with ModernBERT, AIA2 improves WG-F1 over AugGPT by 1.8 points.

5 Related Work

Prior robustness methods address class imbalance or predefined groups, while latent subgroup structure and attribute-driven distribution shifts remain underexplored. Data augmentation, including LLM-based generation, offers a complementary route to improving robustness.

  • Re-weighting and group-robust methods adjust training toward hard classes or predefined groups, but often require observed group information.
  • Latent-subgroup discovery methods identify coherent error slices, hidden subclasses, or concentrated-error regions without explicit subgroup annotations.
  • Traditional and LLM-based augmentation generate paraphrases, mixed samples, or counterfactual examples to improve robustness under distribution shift.

6 Conclusion

AIA2 improves subgroup robustness without explicit attribute annotations by discovering latent slices, repairing local label deficits, and applying targeted LLM augmentation. Across five datasets spanning classification and NER, it improves overall and worst-group performance over competitive baselines.

  • AIA2 discovers latent slices in a joint semantic-predictive space, identifies local label deficits, and uses LLMs for targeted augmentation.
  • Across five text classification and named entity recognition datasets, AIA2 improves overall and worst-group performance over competitive baselines.
  • Generated-data analysis links gains to subgroup-label subsets where augmentation reduces local deficits.
  • Ablations support contributions from slice discovery, gap-aware selection, and LLM generation, with stable results across design choices.

Limitations

AIA2 has computational, representation-quality, generation-quality, and generalizability limitations. Its evaluation covers text classification and NER, while extension to other NLP tasks remains unvalidated.

  • Iterative training, repeated slice discovery, distribution-gap analysis, and LLM augmentation add computational cost that may challenge scaling.
  • Noisy or weakly aligned semantic embeddings and prediction signals can make discovered slices less reliable and affect augmentation quality.
  • LLM-generated examples may contain label noise, hallucinated details, or demographic and topical biases that filtering cannot fully eliminate.
  • The evaluation focuses on text classification and named entity recognition, leaving generalizability to relation extraction and question answering for future validation.

E Stability Analysis

AIA2 shows stable average performance across random seeds while improving subgroup robustness. Worst-group metrics vary more because they depend on difficult subgroups and small sample changes, but AIA2 is more stable than JTT on all worst-group metrics.

  • AIA2 maintains stable average performance across five random seeds while improving subgroup robustness.
  • Macro-F1 and Micro-F1 standard deviations remain below 1.0 point on most datasets.
  • Worst-group metrics show larger variance because they depend on the most difficult subgroup and are sensitive to small sample changes.
  • AIA2 is more stable than JTT on all worst-group metrics, with standard deviations of 1.24 versus 2.44 on WWW2015 and 1.94 versus 2.89 on RE3D.

F Hyperparameter Sensitivity Analysis

AIA2 remains stable across moderate changes to its latent-slice and LLM-augmentation hyperparameters, with the default settings generally competitive.

  • The analysis varies the joint semantic–predictive weighting coefficient, kNN neighbors, Leiden resolution, and generated candidates per seed sample.
  • AIA2 is stable under moderate hyperparameter changes across the evaluated datasets.The sensitivity analysis covers latent slice discovery and LLM-based augmentation on HumAID and CrossNER.
  • Values of α between 0.05 and 0.2 outperform semantic-only slicing across all metrics, while the default α = 0.1 remains competitive.Prediction-only slicing performs worst overall, and larger α values are less consistent.
  • Both datasets remain stable across k ∈ {5, 10, 15, 20}, including the default k = 15.

G Computational Overhead

AIA2 adds computational overhead over lightweight training baselines but is more time-efficient than generic LLM augmentation methods.

  • AIA2 introduces additional overhead over Focal Loss and GroupDRO because of latent slice discovery and LLM-based targeted augmentation.The comparison uses WWW2015 for classification and CrossNER for named entity recognition on an H100 GPU with batch size 16.
  • 15.5% less training time than CB-LLM and 28.2% less than AugGPT are achieved by AIA2 on WWW2015.
  • 10.4% less training time than CB-LLM and 20.2% less than AugGPT are achieved by AIA2 on CrossNER.
  • AIA2 provides a practical cost–benefit tradeoff by remaining more time-efficient than general-purpose LLM augmentation methods.

H Additional Generated Data Analysis

Additional analyses link subgroup-label deficit reduction with subset-level accuracy gains and characterize the retained augmentation volume.

  • ρ = 0.700 (p = 2.93 × 10^-6) indicates a significantly positive association between deficit reduction and subset-level accuracy gain on 35 HumAID subgroup-label subsets.The analysis concerns HumAID and HateXplain and is presented in Figure 6.
  • AIA2 improves subgroup robustness by repairing local subgroup-label deficits rather than relying only on untargeted data expansion.
  • AIA2 retains 1,073 samples on HumAID and 1,643 on CrossNER for training.Table 11 excludes raw candidates rejected during quality control and measures increases relative to the original training-set size.

I Additional Ablations

Ablations support the joint slice representation and seed-selection design, while Leiden provides the most consistent clustering performance across datasets.

  • Representation and selection: The full setting achieves the best results across all metrics on WWW2015 and CrossNER.The ablation compares semantic-only and prediction-only slicing with deficit-only and high-loss-only seed selection.
  • Representation and selection: 3.2–4.3 points on WWW2015 and 1.3–2.0 points on CrossNER improve worst-group performance over single-signal slice variants.
  • Representation and selection: 1.9–2.7 points on WWW2015 and 1.0–1.2 points on CrossNER improve worst-group performance over single-signal selection variants.
  • Clustering methods: Leiden provides the most consistent performance across datasets among the evaluated clustering methods.K-Means is comparable on classification but less effective for irregular error regions; DBSCAN is less stable, and Agglomerative clustering has mixed results.
Loading 2608.30297v1…