Learning Optimal Features via Partial Invariance

TL;DR

This paper introduces a relaxation of Invariant Risk Minimization called partial invariance, which improves robustness to distribution shifts by learning from domain partitions, supported by theoretical analysis and experiments.

Contribution

It proposes partial invariance as a relaxation of IRM, addressing its limitations when causal features are not fully invariant across environments.

Findings

Partial invariance improves model robustness under distribution shifts.

Learning from domain partitions enhances invariant feature learning.

Experimental results validate the theoretical advantages of partial invariance.

Abstract

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.