Reliability-Aware Environment Discovery: Leveraging Feature Entanglement for Subpopulation Robustness
Abstract
Machine learning models often fail under subpopulation shift, where latent subgroup distributions differ between training and test environments despite stable within-group relationships between features and labels. While empirical risk minimization achieves high average accuracy, it frequently performs poorly on minority subpopulations. Existing group-robust methods address this issue by optimizing worst-group risk, but typically require subgroup annotations or rely on error-based environment discovery. These approaches implicitly assume that failures are driven by dominant spurious correlations, leading to coherent error patterns. However, this assumption breaks down when spurious and invariant features are entangled, yielding diffuse and heterogeneous failures that are not well captured by error frequency alone. We propose Reliability-Aware Environment Discovery (RAE), an annotation-free framework that incorporates prediction reliability as an auxiliary signal for discovering vulnerable environments. RAE quantifies reliability using split conformal nonconformity scores, enabling the identification of regions with insufficient or contradictory predictive evidence. Then, define group label by combining prediction errors with reliability signals and apply group-robust optimization over the discovered partition. Experiments on vision and language benchmarks demonstrate that RAE improves both average and worst-group performance, with discovered environments closely approximating oracle group supervision.