Invariant Feature Subspace Recovery for Multi-Class Classification

Gargi Balasubramaniam, Haoxiang Wang, Han Zhao

Keywords: Domain Generalization, Invariant Feature Learning, Out of distribution (OOD) Generalization, Spurious Correlation Robustness

TL;DR: We propose ISR-Multiclass, a provable domain generalization algorithm that efficiently recovers the invariant-feature subspace in the setting of multi-class classification.

Abstract: Domain generalization aims to learn a model over multiple training environments to generalize to unseen environments. Recently, Wang et al. [2022] proposed Invariant-feature Subspace Recovery (ISR), a domain generalization algorithm that uses the means of class-conditional data distributions to provably identify the invariant-feature subspace under a given causal model. However, due to the specific assumptions of the causal model, the original ISR algorithm is conditioned on a single class only, without utilizing information from the rest of the classes. In this work, we consider the setting of multi-class classification under a more general causal model, and propose an extension of the ISR algorithm, called ISR-Multiclass. This proposed algorithm can provably recover the invariant-feature subspace with $\lceil d_{spu}/k \rceil + 1$ environments, where $d_{spu}$ is the number of spurious features and $k$ is the number of classes. Empirically, we first examine ISR-Multiclass in a synthetic dataset, and demonstrate its superiority over the original ISR in the multi-class setting. Furthermore, we conduct experiments in Multiclass Coloured MNIST, a semi-synthetic dataset with strong spurious correlations, and show that ISR-Multiclass can significantly improve the robustness of neural nets trained by various methods (e.g., ERM and IRM) against spurious correlations.