Clustered Invariant Risk Minimization

Abstract

This study extends the problem settings of Invariant Risk Minimization(IRM) for Out-of-Distribution generalization problems to unknown clustered environments settings. In this scenario, where a given set of environments exhibits an unknown clustered structure, our objective is to identify a single invariant feature extractor and per-cluster regressors (or classifiers) built on top of the feature extractor. To achieve this, we propose a new framework called Clustered IRM for simultaneously identifying the cluster structure and the invariant features. Our theoretical analysis demonstrates that the required number of training environments for such identification is only $O(d_\mathrm{sp} + K^2)$, where $d_\mathrm{sp}$ represents the dimensionality of the spurious features, and $K$ is the number of clusters. Numerical experiments validate the effectiveness of our proposed framework.