Robustness Beyond Known Groups with Low-rank Adaptation

Deep learning models trained to optimize average accuracy often exhibit systematic failures on particular subpopulations. In real world settings, the subpopulations most affected by such disparities are frequently unlabeled or unknown, thereby motivating the development of methods that are performant on sensitive subgroups without being pre-specified. However, existing group-robust methods typically assume prior knowledge of relevant subgroups, using group annotations for training or model selection. We propose Low-rank Error Informed Adaptation (LEIA), a simple two-stage method that improves group robustness by identifying a low-dimensional subspace in the representation space where model errors concentrate. LEIA restricts adaptation to this error-informed subspace via a low-rank adjustment to the classifier logits, directly targeting latent failure modes without modifying the backbone or requiring group labels. Using five real-world datasets, we analyze group robustness under three settings: (1) truly no knowledge of subgroup relevance, (2) partial knowledge of subgroup relevance, and (3) full knowledge of subgroup relevance. Across all settings, LEIA consistently improves worst-group performance while remaining fast, parameter-efficient, and robust to hyperparameter choice.

💡 Research Summary

The paper tackles a fundamental challenge in modern machine learning: models trained to maximize average accuracy often fail dramatically on subpopulations that are not identified during training. In high‑stakes domains such as healthcare, hiring, or hate‑speech detection, these vulnerable groups are frequently unknown or lack reliable annotations, making traditional group‑robust methods impractical. Existing approaches like Group DRO, JTT, or AFR either require explicit group labels for training and validation or assume that all relevant groups are known a priori. The authors formally demonstrate (Proposition 3.3) that optimizing over a known set of groups can actually degrade performance on latent, unannotated groups, highlighting the need for methods that do not depend on group information.

To address this gap, the authors introduce Low‑rank Error Informed Adaptation (LEIA), a two‑stage procedure that improves worst‑group performance without any group labels. In Stage 1, a standard Empirical Risk Minimization (ERM) model is trained on a large portion of the data, producing a feature extractor (e_{\psi}) and a classifier head (c_{\phi}). After convergence, the feature extractor is frozen. Stage 2 focuses on the remaining data subset: each example receives a weight (\mu_i \propto \exp(\gamma \cdot \ell)) where (\ell) is the cross‑entropy loss, thereby emphasizing high‑loss (i.e., error‑prone) samples. Using these weights, the authors compute an error‑weighted covariance matrix (\Sigma_{\text{err}}) of the frozen representations and extract its top‑(k) eigenvectors (V_k). This low‑dimensional subspace captures directions in representation space where errors concentrate.

LEIA then learns a single low‑rank matrix (A \in \mathbb{R}^{k \times C}) that adds a correction term (A^\top V_k^\top e_{\psi}(x)) directly to the logits. The correction is trained by minimizing a weighted cross‑entropy loss over the same subset, with only (A) being updated. Consequently, the method adds a negligible number of parameters (often <0.1 % of the original model) and incurs minimal computational overhead because the backbone remains frozen.

The authors evaluate LEIA on five real‑world benchmarks spanning vision (Waterbirds, CelebA) and language (MultiNLI, CivilComments, etc.) under three knowledge regimes: (1) no group information, (2) partial group information, and (3) full group information. Across all settings, LEIA consistently improves worst‑group accuracy relative to 12 strong baselines, achieving absolute gains of 2–5 percentage points. Notably, in the fully unlabeled scenario, LEIA outperforms Group DRO and JTT despite having no access to group labels for hyper‑parameter tuning or early stopping. The method also demonstrates strong parameter efficiency: the only learnable component is the low‑rank matrix (A), and training time is reduced by roughly 30 % compared to full‑model fine‑tuning approaches.

Ablation studies explore the impact of the rank (k) and the sharpness parameter (\gamma). The cumulative explained variance (CEV) analysis shows that a small number of eigenvectors (typically 5–20) capture 50–90 % of the error variance, confirming the low‑rank nature of the error structure. Larger (\gamma) values focus the adaptation on the most mis‑classified examples, yielding more pronounced improvements but with diminishing returns beyond a certain point. The authors also propose a data‑driven heuristic for selecting (k) based on the proportion of variance explained, which reduces the need for extensive hyper‑parameter sweeps.

Beyond quantitative results, the paper provides visualizations that illustrate how the identified error subspace aligns with mis‑classified points from the worst‑performing group, and how the low‑rank correction shifts the decision boundary to rectify these errors without harming overall accuracy. Theoretical discussion, detailed proofs (Appendix B), and extensive supplementary material (Appendix F) situate LEIA within the broader literature on invariant risk minimization, feature reweighting, and low‑rank adaptation (e.g., LoRA).

In summary, LEIA offers a practical, theoretically grounded, and computationally light solution for enhancing group robustness when subgroup information is unknown, partially known, or fully known. By leveraging the geometric structure of model errors rather than explicit group annotations, it bridges a critical gap between fairness‑oriented research and real‑world deployment constraints. Future directions suggested include extending the error‑subspace extraction to multi‑stage or hierarchical adaptations and integrating unsupervised clustering to discover even richer latent group structures.