Out-of-Distribution Optimality of Invariant Risk Minimization

TL;DR

This paper provides a theoretical foundation showing that invariant risk minimization (IRM) solutions can minimize out-of-distribution risk under certain conditions, addressing a key gap in understanding IRM's guarantees.

Contribution

It rigorously proves that IRM's bi-level optimization minimizes out-of-distribution risk given specific distributional and dimensional conditions.

Findings

IRM solutions minimize o.o.d. risk under certain conditions

Provides sufficient conditions on data distributions for IRM effectiveness

Establishes theoretical guarantees for IRM's optimality

Abstract

Deep Neural Networks often inherit spurious correlations embedded in training data and hence may fail to generalize to unseen domains, which have different distributions from the domain to provide training data. M. Arjovsky et al. (2019) introduced the concept out-of-distribution (o.o.d.) risk, which is the maximum risk among all domains, and formulated the issue caused by spurious correlations as a minimization problem of the o.o.d. risk. Invariant Risk Minimization (IRM) is considered to be a promising approach to minimize the o.o.d. risk: IRM estimates a minimum of the o.o.d. risk by solving a bi-level optimization problem. While IRM has attracted considerable attention with empirical success, it comes with few theoretical guarantees. Especially, a solid theoretical guarantee that the bi-level optimization problem gives the minimum of the o.o.d. risk has not yet been established. Aiming at providing a theoretical justification for IRM, this paper rigorously proves that a solution to the bi-level optimization problem minimizes the o.o.d. risk under certain conditions. The result also provides sufficient conditions on distributions providing training data and on a dimension of feature space for the bi-leveled optimization problem to minimize the o.o.d. risk.