SAFE: Stable Feature Extraction without Environment Labels

We study the problem of stable (or invariant) feature extraction for OOD generalization, when environment labels are unknown. Prior works extract stable features by first inferring pseudo environment labels, before applying invariant learning methods like invariant risk minimization (IRM). These methods are highly sensitive to hyper-parameters and model selection strategies. Moreover, recent work shows that it is not always possible to identify stable features. In this paper, we present sufficient conditions under which stable features can be directly identified without environment labels. Using these conditions, we provide a practical algorithm called StAble Feature Extraction ($\textbf{SAFE}$), which selects stable features without inferring environment labels. We show that SAFE accurately removes spurious and selects stable features on synthetic as well as a diverse range of real-world datasets, improving the OOD performance and calibration of ERM as well as prior invariant learning algorithms. Our work highlights the inefficacy of current invariant learning methods, and calls for more attention to the identifiability problem of stable features.