Invariant Risk Minimization Is A Total Variation Model

TL;DR

This paper reveals that invariant risk minimization (IRM) is fundamentally a total variation model and introduces a new TV-$ ext{l}_1$ based IRM framework that enhances robustness and generalization in machine learning.

Contribution

The paper provides a mathematical explanation of IRM as a total variation model and proposes a novel TV-$ ext{l}_1$ IRM framework for improved invariant feature learning.

Findings

IRM is essentially a total variation model based on the $L^2$ norm.

The proposed TV-$ ext{l}_1$ IRM framework expands function classes and improves robustness.

Experimental results show competitive performance in benchmark scenarios.

Abstract

Invariant risk minimization (IRM) is an arising approach to generalize invariant features to different environments in machine learning. While most related works focus on new IRM settings or new application scenarios, the mathematical essence of IRM remains to be properly explained. We verify that IRM is essentially a total variation based on $L^2$ norm (TV-$\ell_2$) of the learning risk with respect to the classifier variable. Moreover, we propose a novel IRM framework based on the TV-$\ell_1$ model. It not only expands the classes of functions that can be used as the learning risk and the feature extractor, but also has robust performance in denoising and invariant feature preservation based on the coarea formula. We also illustrate some requirements for IRM-TV-$\ell_1$ to achieve out-of-distribution generalization. Experimental results show that the proposed framework achieves competitive performance in several benchmark machine learning scenarios.