Learning Optimal Features via Partial Invariance
This work addresses robustness in machine learning models for real-world applications, but it is incremental as it builds on the existing IRM framework.
The paper tackles the problem of learning robust models under distribution shifts by showing that Invariant Risk Minimization (IRM) can be suboptimal when its invariance assumption is not fully satisfied, and proposes a relaxation via partial invariance to improve performance. Experiments on language and image data verify the conclusions.
Learning models that are robust to distribution shifts is a key concern in the context of their real-life applicability. Invariant Risk Minimization (IRM) is a popular framework that aims to learn robust models from multiple environments. The success of IRM requires an important assumption: the underlying causal mechanisms/features remain invariant across environments. When not satisfied, we show that IRM can over-constrain the predictor and to remedy this, we propose a relaxation via $\textit{partial invariance}$. In this work, we theoretically highlight the sub-optimality of IRM and then demonstrate how learning from a partition of training domains can help improve invariant models. Several experiments, conducted both in linear settings as well as with deep neural networks on tasks over both language and image data, allow us to verify our conclusions.