# Invariant Risk Minimization Is A Total Variation Model

Zhao-Rong Lai  Weiwen Wang

###### 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.

Invariant risk minimization, total variation model, coarea formula, out-of-
distribution generalization

##  1 Introduction

Many machine learning tasks can be reduced to minimizing the learning risk,
which is properly designed according to the task. Specifically, the learning
risk is empirically minimized on the training set, but the learned model
should be used on the unseen data. If there are significant differences
between the features of the training and unseen data, the performance of the
learning system may deteriorate (Recht et al., 2019; Arjovsky et al., 2020;
Lin et al., 2022). For example, if a convolutional neural network (CNN) is
trained with pictures of cows in grass and camels in deserts, it may fail to
classify easy samples of cows in deserts. This is because the CNN minimizes
the training error by classifying grass and deserts as cows and camels,
respectively.

The key to solving this misclassification problem is to distinguish between
the invariant features (animal shapes) and the spurious features (landscapes)
which lead to distributional shift. Invariant Risk Minimization (IRM,
(Arjovsky et al., 2020)) emerges as a learning paradigm that estimates
nonlinear, invariant, and causal predictors from multiple training
environments for such out-of-distribution (OOD) generalization issue. It
introduces a gradient norm penalty that measures the optimality of the dummy
classifier at each environment. Different variants of IRM have been proposed
since then, such as Risk Extrapolation (REx, Krueger et al. 2021), SparseIRM
(Zhou et al., 2022), jointly-learning with auxiliary information (ZIN, Lin et
al. 2022), and invariant feature learning through independent variables (TIVA,
Tan et al. 2023). These works improve the robustness and extendability of IRM.

In this paper, we investigate the mathematical essence of IRM and find that it
can be absorbed in a traditional and widely-used operator in mathematics:
total variation (TV). TV has long been used in various fields of mathematics
and engineering, like optimal control, data transmission, sensing and
denoising. For example, TV based on $\ell_{1}$ norm (TV-$\ell_{1}$, (Rudin et
al., 1992)) is exploited to develop noise removal algorithms in signal
processing. TV-$\ell_{1}$ can also be used in image restoration and processing
(Chan et al., 2006; Chen et al., 2006). In other aspects, TV-$\ell_{1}$ is
also a tractable approach to data-driven scale selection and function
approximation (Chan & Esedoglu, 2005).

Based on this motivation, we reveal some useful mathematical properties of TV-
based IRM models that can be exploited in machine learning scenarios. 1. We
formulate and verify that the existing IRM framework is essentially a
TV-$\ell_{2}$ model. 2. We propose a novel IRM framework based on the
TV-$\ell_{1}$ model (IRM-TV-$\ell_{1}$). It has two advantages: 1) The set of
TV-$\ell_{1}$ integrable functions is larger than that of TV-$\ell_{2}$
integrable functions when the measure is finite (this is the case in IRM-
TV-$\ell_{2}$), hence more classes of functions are allowed as the learning
risk and the feature extractor in IRM-TV-$\ell_{1}$. 2) TV-$\ell_{1}$ has
robust performance in denoising and preserving sharp features based on the
coarea formula. This property helps to shape a blocky (piecewise-constant)
learning risk that is more robust to the environment change. 3. We investigate
some requirements for IRM-TV-$\ell_{1}$ to achieve OOD generalization, such as
a more flexible penalty parameter, extendability of the training environment
set, and accuracy of the measure. These findings help to explain why IRM is
effective theoretically.

##  2 Preliminaries and Related Works

Note that there are different formulations and interpretations for both IRM
and TV models. For the purpose of understanding the proposed approach, we
adopt the most convenient formulations without loss of generality.

###  2.1 Invariant Risk Minimization

Suppose we have a training data set of $n$ samples
$\mathcal{D}:=\\{(x_{i},y_{i})\in\mathcal{X}\times\mathcal{Y}\\}_{i=1}^{n}$,
where $\mathcal{X}$ and $\mathcal{Y}$ denote the input space and the output
space, respectively. These samples are collected from some training
environments in the set $\mathcal{E}_{tr}$ (e.g., grass or deserts). Note that
we take a general scenario where environment partition is absent, thus the
true environment label $e$ for a sample $(x_{i},y_{i})$ is unknown. The
machine learning task is to learn a predictor
$\mathcal{X}\rightarrow\mathcal{Y}$ such that some predefined risk (or error)
metric $R$ can be minimized on the training data $\mathcal{D}$. In the IRM
framework (Arjovsky et al., 2020), this predictor is composed of two operators
$w\circ\Phi$, where $\Phi:\mathcal{X}\rightarrow\mathcal{H}$ is able to
extract invariant features and $w:\mathcal{H}\rightarrow\mathcal{Y}$ is a
classifier.

Suppose we use the loss function
$\mathcal{L}:\mathcal{Y}\times\mathcal{Y}\times\mathcal{E}_{tr}\rightarrow\mathbb{R}$
to compute the prediction error, where the environment variable becomes an
argument of $\mathcal{L}$. The empirical risk of the training samples in the
environment $e$ is computed by

$R(w\circ\Phi,e):=\frac{1}{n}\sum_{i=1}^{n}\mathcal{L}(w\circ\Phi(x_{i}),y_{i},e),$
(1)

which is the mean loss between the predicted value $w\circ\Phi(x_{i})$ and the
true value $y_{i}$ in the environment $e$. The notations $x_{i}$ and $y_{i}$
can be omitted in $R$ for concise expressions, as they are constants. The
actual parameters that we should learn are $w$, $\Phi$, and $e$.

The original IRM (Arjovsky et al., 2020) can be defined by

$\displaystyle\min_{w,\Phi}\ $
$\displaystyle\sum_{e\in\mathcal{E}_{tr}}R(w\circ\Phi,e)$ s. t. $\displaystyle
w\in\mathop{\rm arg\,min}_{\tilde{w}}R(\tilde{w}\circ\Phi,e),\quad\forall
e\in\mathcal{E}_{tr}.$ (IRM)

It aims to learn $w$ and $\Phi$ simultaneously by minimizing the total risk
and forcing $w$ to be a uniform minimizer of the risk under each environment.
However, since it is a challenging bi-level optimization problem, a practical
surrogate can be used instead (Arjovsky et al., 2020):

$\min_{\Phi}\
\sum_{e\in\mathcal{E}_{tr}}\left\\{R(1\circ\Phi,e)+\lambda\|\nabla_{w}\arrowvert_{w=1}R(w\circ\Phi,e)\|_{2}^{2}\right\\},$
(IRMv1)

where $\Phi$ becomes the entire invariant predictor and the classifier $w$ is
fixed to be a scalar $1$. The penalty parameter $\lambda$ is non-negative, and
the $\ell_{2}$ norm is denoted by $\|\cdot\|_{2}$. The minimization in the
constraint of (2.1) is converted into a gradient norm penalty in the objective
of (IRMv1).

The REx (Krueger et al., 2021) approach extrapolates the training risks to
their affine combinations. The Variance-REx (V-REx) is a simple, stable and
effective form of REx that takes the variance of training risks in different
environments as the regularizer:

$\min_{\Phi}\
\sum_{e\in\mathcal{E}_{tr}}R(1\circ\Phi,e)+\lambda\mathbb{V}(\\{R(1\circ\Phi,e)\\}_{e\in\mathcal{E}_{tr}}),$
(V-REx)

where $\mathbb{V}(\cdot)$ denotes the variance operator. The scalar $w=1$ is
retained to be consistent with (IRMv1).

As for other extensions of IRM, ZIN (Lin et al., 2022) further exploits
additional auxiliary information $\\{z_{i}\in\mathcal{Z}\\}_{i=1}^{n}$ to
simultaneously learn environment partition and invariant representation. It
assumes that the training environment space $\mathcal{E}_{tr}$ is a convex
hull of $E$ linearly-independent fundamental environments (basis). Then each
environment in $\mathcal{E}_{tr}$ is isomorphic to an element in the
$E$-dimensional simplex

$\displaystyle\Delta^{E}:=\\{v\in\mathbb{R}^{E}:v\geqslant 0_{(E)}\
\text{and}\ v\cdot 1_{(E)}=1\\},$ (2)

where $0_{(E)}$ and $1_{(E)}$ denote $E$-dimensional vectors of $0$ and $1$,
respectively. The inner product of two vectors is denoted by $\cdot$. ZIN aims
to learn a mapping $\rho:\mathcal{Z}\rightarrow\Delta^{E}$ that directly
serves as the weights for different environments:

$\displaystyle
R(w\circ\Phi,\rho):=\frac{1}{n}\sum_{i=1}^{n}\tilde{\mathcal{L}}(w\circ\Phi(x_{i}),y_{i})\rho(z_{i})\in\mathbb{R}^{E},$
$\displaystyle\min_{w,\Phi}\
\left\\{R(w{\circ}\Phi,\frac{1}{E}1_{(E)}){\cdot}1_{(E)}{+}\lambda\max_{\rho}\|\nabla_{w}R(w{\circ}\Phi,\rho)\|_{2}^{2}\right\\}.$
(ZIN)

Compared with (1), (2.1) decomposes the environment argument into a convex
combination of $E$ components. The treatment of the environment lies in
learning $\rho$. $\frac{1}{E}1_{(E)}$ denotes the averaged environment, and
$R(w\circ\Phi,\frac{1}{E}1_{(E)})\cdot 1_{(E)}$ measures the holistic
empirical loss across all environments. The variation
$\nabla_{w}R(w\circ\Phi,\rho)\in\mathbb{R}^{E\times d_{\mathcal{H}}}$ becomes
a gradient matrix, and the matrix $\ell_{2}$ norm is used as a penalty. The
penalty is maximized w.r.t. $\rho$ to identify the worst environment that
causes the largest absolute variation. Then the whole risk and penalty are
minimized w.r.t. $(w,\Phi)$.

TIVA (Tan et al., 2023) directly uses the auxiliary information as input
variables
$\\{\tilde{x}_{i}:=(x_{i},z_{i})\in\mathcal{X}\times\mathcal{Z}\\}_{i=1}^{n}$
instead of sending it to the environment predictor $\rho$. It decomposes
$\tilde{x}:=(\tilde{x}\odot u^{+})+(\tilde{x}\odot(1-u^{+}))$ with a gate
function $u^{+}\in[0,1]^{d_{\mathcal{X}\times\mathcal{Z}}}$ to identify
environmental-related features. The risk metric of TIVA is defined by

$\displaystyle
R(w{\circ}\Phi,\rho):=\frac{1}{n}\sum_{i=1}^{n}\tilde{\mathcal{L}}(w{\circ}\Phi(\tilde{x}{\odot}u^{+}),y_{i})\rho(\tilde{x}{\odot}(1{-}u^{+})).$
(TIVA)

The optimization model is similar to (2.1).

###  2.2 Total Variation

Total variation (TV, Rudin et al. 1992) is a useful operator that
characterizes the varying property of a function. Let
$\Omega\subseteq\mathbb{R}^{d}$ be an open set and $L^{1}(\Omega)$ be the
space of functions defined on $\Omega$ with bounded Lebesgue integrals. Then
for any function $f\in L^{1}(\Omega)$, its TV can be defined by (Chan et al.,
2006)

$\displaystyle\int_{\Omega}|\nabla f|:=\sup\\{\int_{\Omega}$ $\displaystyle
f(x)\mathrm{div}g(x)\,\mathrm{d}x:\ g\in C_{c}^{1}(\Omega,\mathbb{R}^{d}),$
$\displaystyle\|g\|_{L^{\infty}(\Omega)}\leqslant 1\\},$ (TV)

where $g$ is a differentiable vector function with compact support contained
in $\Omega$ and essential supremum no larger than $1$. $\mathrm{div}g$ denotes
the divergence of $g$. In brief, the TV operator measures the local change of
$f$ along all dimensions of $x$, then integrates all these local changes
throughout the whole domain $\Omega$. Note that $f$ is not necessarily
differentiable and the notation $|\nabla f|$ is symbolic. But if $f$ is
differentiable, $|\nabla f|$ truly becomes the modulus (i.e., $\ell_{2}$ norm)
of the gradient.

Only functions with bounded TV are meaningful in this realm, leading to the
bounded variation (BV) space $\\{f\in L^{1}(\Omega):\int_{\Omega}|\nabla
f|<\infty\\}$. To simplify expressions, the term TV function actually means
the BV function in this paper. In general, a TV model aims to minimize TV by
default without explicitly stating it, because it makes little sense to
maximize TV in practice.

The TV operator is known for preserving sharp discontinuities while removing
noise and other unwanted fine scale detail of $f$. This property is based on
the coarea formula (Chan et al., 2006)

$\displaystyle\int_{\Omega}|\nabla
f|:=\int_{-\infty}^{\infty}\int_{f^{-1}(\gamma)}\,\mathrm{d}s\,\mathrm{d}\gamma,$
(3)

where $f^{-1}(\gamma):=\\{x\in\Omega:f(x)=\gamma\\}$ denotes the level set
(preimage) of $f$ at $\gamma$. TV actually integrates along all contours of
$f^{-1}(\gamma)$ for all $\gamma$ where the differential $\,\mathrm{d}\gamma$
exists. Hence if $f$ is more blocky (piecewise-constant), it will have a
smaller TV. This property is exploited in many signal recovery applications
via the following TV-$\ell_{1}$ model (Rudin et al., 1992)

$\displaystyle\inf_{f\in L^{2}(\Omega)}\left\\{\int_{\Omega}|\nabla
f|+\lambda\int_{\Omega}(f-\tilde{f})^{2}\,\mathrm{d}x\right\\},$
(TV-$\ell_{1}$)

where $\tilde{f}\in L^{2}(\Omega)$ is the ground-truth signal to be recovered
and $\lambda$ is the approximating accuracy parameter. The objective of this
model is to preserve sharp discontinuities in the approximation $f$ while
leaving noise (especially Gaussian noise) and other unwanted fine scale detail
in the residue $(f-\tilde{f})$.

If $f$ satisfies the variation condition $\int_{\Omega}|\nabla f|^{2}<\infty$,
then the following TV-$\ell_{2}$ model (Mumford & Shah, 1985) is feasible:

$\displaystyle\inf_{f\in L^{2}(\Omega)}\left\\{\int_{\Omega}|\nabla
f|^{2}+\lambda\int_{\Omega}(f-\tilde{f})^{2}\,\mathrm{d}x\right\\}.$
(TV-$\ell_{2}$)

In fact, TV-$\ell_{2}$ appears earlier in history than TV-$\ell_{1}$. However,
the square-integral $\int_{\Omega}|\nabla f|^{2}$ does not have the coarea
formula (3). Therefore, TV-$\ell_{2}$ generally does not produce a blocky
recovered signal $f$.

##  3 Methodology

Motivated by the properties of TV, we aim to develop a risk metric that is
blocky (piecewise-constant) w.r.t. the environment, so that it can be better
generalized to different environments. First, we establish the theoretical
framework for IRM to be a TV-$\ell_{2}$ model.

###  3.1 Conditions for IRM to be TV-$\ell_{2}$

A necessary but not sufficient condition for the classifier $w$ to achieve the
minimums in the constraints of (2.1) is

$\displaystyle\nabla_{w}R(w\circ\Phi,e)=0,\quad\forall e\in\mathcal{E}_{tr}.$
(4)

This leads to the surrogate model (IRMv1) which penalizes the gradient norms
over all the environments. However, the risk function $R(w\circ\Phi,e)$
usually has a complex nonconvex geometric structure w.r.t. $w$ in deep
learning. Thus (4) may achieve a maximum instead of a minimum, and it is
intractable to turn (IRMv1) back to (2.1).

From a different perspective, we interpret (IRMv1) as a TV-$\ell_{2}$ model in
two ways: 1. Consider the environment features as noise and try to remove it.
2. Recover the useful signal (e.g., animal shapes) from the background
environment.

To begin with, we turn $w$ back to an argument in the minimization of (IRMv1).
This is reasonable because the learning of $w$ and $\Phi$ is a dynamic
process.

$\displaystyle\min_{w,\Phi}\ $
$\displaystyle\sum_{e\in\mathcal{E}_{tr}}\left\\{R(w\circ\Phi,e)+\lambda\|\nabla_{w}R(w\circ\Phi,e)\|_{2}^{2}\right\\}.$
(5)

We also need the following successive conditions.

Condition 1. There exists a measure $\mu$ for $e$ on
$(\mathcal{E}_{tr},\mathscr{F}_{tr})$ where $\mathscr{F}_{tr}$ denotes the
$\sigma$-algebra on $\mathcal{E}_{tr}$ that includes all the environment
combinations in the training set.

Condition 2. Under this measure $\mu$, the feature extractor $\Phi$ is
uncorrelated to the environment variable $e$. Moreover, the correlation of the
risk metric $R$ to $e$ only lies on the classifier $w$, i.e.,
$R(w\circ\Phi,e)\,\mathrm{d}\mu\equiv R(w(e)\circ\Phi)\,\mathrm{d}\mu$.

Condition 3. $w$ is a measurable function of $e$ in the sense that $w$ is
parameterized (either fully or partly) by $e$. In other words,
$w(\cdot)\circ\Phi:\mathcal{E}_{tr}\times\mathcal{H}\rightarrow\mathcal{Y}$ is
measurable.

Condition 4. $R(w(e)\circ\Phi)\in
L^{1}(\mathcal{E}_{tr},\mathscr{F}_{tr},\mu)$ and
$|\nabla_{w}R(w(e)\circ\Phi)|\in
L^{2}(\mathcal{E}_{tr},\mathscr{F}_{tr},\mu)$.
$L^{1}(\mathcal{E}_{tr},\mathscr{F}_{tr},\mu)$ and
$L^{2}(\mathcal{E}_{tr},\mathscr{F}_{tr},\mu)$ denote the function spaces with
integrable and square-integrable functions under measure $\mu$ defined on
$(\mathcal{E}_{tr},\mathscr{F}_{tr})$, respectively.

Condition 1 implies that we are able to measure the magnitude of the
environments we encounter in the training set. For instance, (5) actually uses
a counting measure on $\mathcal{E}_{tr}$ such that for any
$e\in\mathcal{E}_{tr}$, $\,\mathrm{d}\mu(\\{e\\})=1$. Condition 2 is weaker
than $R(w\circ\Phi,e)\equiv R(w(e)\circ\Phi)$ and indicates that the
environment influences only $w$ but not $\Phi$ under an infinitesimal scale of
environment change $\,\mathrm{d}\mu$. Then $w$ passes the environment
representation to $R$. This is reasonable because $\Phi$ is the invariant
feature extractor that we intend to learn. Condition 3 enables us to represent
and measure $w$ by $e$. Finally, Condition 4 is fundamental to ensure
boundedness when we compute the total risk and its TV.

######  Theorem 3.1.

A risk metric $R(w\circ\Phi,e)$ that satisfies Conditions 1$\sim$4 has the
following well-defined finite integral and TV-$\ell_{2}$ form w.r.t. $w$:

$\displaystyle\int_{\Omega}R(w\circ\Phi)\,\mathrm{d}\nu\quad\text{and}\quad\int_{\Omega}|\nabla_{w}R(w\circ\Phi)|^{2}\,\mathrm{d}\nu,$
(6)

where $\Omega=w(\mathcal{E}_{tr})$ is the image of $\mathcal{E}_{tr}$ under
mapping $w$, and $\nu$ is the measure for $w$ induced by $\mu$. If $\mu$ is a
probability measure, then the above integrals become mathematical expectations

$\displaystyle\mathbb{E}_{w}[R(w\circ\Phi)]\quad\text{and}\quad\mathbb{E}_{w}[|\nabla_{w}R(w\circ\Phi)|^{2}].$
(7)

The following IRM-TV-$\ell_{2}$ model is also well-defined:

$\displaystyle\min_{\Phi}\ $
$\displaystyle\mathbb{E}_{w}[R(w\circ\Phi)+\lambda|\nabla_{w}R(w\circ\Phi)|^{2}].$
(IRM-TV-$\ell_{2}$)

The proof is presented in Appendix A.1. Note that only $\Phi$ is left in the
minimizing argument because $w$ has been taken expectation in the objective
function. This reveals the mathematical essence of why $w$ can be set as a
dummy scalar classifier $1$ in (IRMv1).

It follows from Theorem 3.1 and multiplying (5) by $1/|\mathcal{E}_{tr}|$ that
the following corollary can be obtained.

######  Corollary 3.2.

Let $\mathcal{E}_{tr}$ be a finite discrete environment training set with a
discrete uniform probability measure $\mu$. Then given Conditions 1$\sim$4,
(5) can be generalized to (IRM-TV-$\ell_{2}$).

By slightly changing Condition 4, (V-REx) can also be generalized like Theorem
3.1.

Condition 4a. $R(\Phi,e)\in L^{2}(\mathcal{E}_{tr},\mathscr{F}_{tr},\mu)$ with
$\mu$ being a probability measure.

######  Proposition 3.3.

Given Conditions 1 and 4a, (V-REx) can be generalized to the following
V-REx-$\ell_{2}$ model:

$\displaystyle\min_{\Phi}\ $
$\displaystyle\mathbb{E}_{e}[R(\Phi,e)]{+}\lambda\mathbb{E}_{e}[(R(\Phi,e){-}\mathbb{E}_{e}[R(\Phi,e)])^{2}].$
(V-REx-$\ell_{2}$)

The proof is presented in Appendix A.2. In (V-REx-$\ell_{2}$),
$\mathcal{E}_{tr}$ is not necessarily a finite discrete set and $\mu$ is not
necessarily a uniform probability measure. Besides, (V-REx-$\ell_{2}$) does
not require a classifier $w$ to convey environment information (Conditions 2
and 3). This proposition reveals the mathematical essence of why V-REx can
extrapolate risk to a larger region (see Appendix A.2). (V-REx-$\ell_{2}$) can
be considered as a variant of TV-$\ell_{2}$ model, as variance is a similar
tool to TV that assesses the deviation of a function.

To incorporate ZIN and TIVA into our TV-$\ell_{2}$ framework, we present the
following theorem.

######  Theorem 3.4.

The environment learner $\rho$ is actually a probability measure satisfying
Condition 1. Given Conditions 2$\sim$4, ZIN and TIVA can be generalized to the
following Minimax-TV-$\ell_{2}$ model

$\displaystyle\min_{\Phi}\ $
$\displaystyle\\{\mathbb{E}_{w\leftarrow\frac{1_{(E)}}{E}}[R(w{\circ}\Phi)]{+}\lambda\max_{\rho}\mathbb{E}_{w\leftarrow\rho}[|\nabla_{w}R(w{\circ}\Phi)|^{2}]\\},$
(Minimax-TV-$\ell_{2}$)

where $\mathbb{E}_{w\leftarrow\rho}$ denotes the mathematical expectation
w.r.t. $w$ whose measure $\nu$ is induced by $\rho$, and $\frac{1_{(E)}}{E}$
denotes the uniform probability measure for $e$.

The proof is presented in Appendix A.3. This theorem indicates that this
minimax scheme actually selects a probability measure $\rho$ that can maximize
the risk variation (the worst case), then minimizes this worst-case total
objective by $\Phi$.

###  3.2 IRM-TV-$\ell_{1}$

In last section, we verify that some typical IRM models and related extensions
are essentially TV-$\ell_{2}$ models. However, IRM-TV-$\ell_{2}$ does not have
the coarea formula (Chan et al., 2006) that provides a geometric nature of
sharp discontinuity preservation, despite having the denoising property. As a
result, the processed learning risk $R$ may still be environment-sensitive.

To address this issue, we establish a novel IRM-TV-$\ell_{1}$ framework that
enhances the robustness of learning risk to environment changes based on the
coarea formula. Since $L^{2}(\Omega,\mathscr{F}_{\Omega},\nu)\subsetneqq
L^{1}(\Omega,\mathscr{F}_{\Omega},\nu)$ when $\nu(\Omega)<\infty$ (this is
usually the case especially when $\nu$ is a probability measure, including
IRM-TV-$\ell_{2}$ and Minimax-TV-$\ell_{2}$), IRM-TV-$\ell_{1}$ allows broader
classes of functions to be the learning risk $R$ and the feature extractor
$\Phi$, which enables us to address more complicated tasks and explore more
feature extractors, respectively. To achieve this, we can relax Condition 4 as
the following Condition 4b. It is a necessary but not sufficient condition for
Condition 4 when $\mu(E_{tr})<\infty$.

Condition 4b. $R(w(e)\circ\Phi),|\nabla_{w}R(w(e)\circ\Phi)|\in
L^{1}(\mathcal{E}_{tr},\mathscr{F}_{tr},\mu)$.

######  Theorem 3.5.

$\Omega$, $\nu$, $\mathbb{E}_{w}$ and $\rho$ are the same as those in Theorems
3.1 and 3.4. A risk metric $R(w\circ\Phi,e)$ satisfying Conditions 1$\sim$3
and 4b has the following well-defined integral, TV-$\ell_{1}$ form and
mathematical expectations w.r.t. $w$:

$\displaystyle\int_{\Omega}R(w\circ\Phi)\,\mathrm{d}\nu\quad\text{and}\quad\int_{\Omega}|\nabla_{w}R(w\circ\Phi)|\,\mathrm{d}\nu,$
(8)
$\displaystyle\mathbb{E}_{w}[R(w\circ\Phi)]\quad\text{and}\quad\mathbb{E}_{w}[|\nabla_{w}R(w\circ\Phi)|].$
(9)

The following IRM-TV-$\ell_{1}$ and Minimax-TV-$\ell_{1}$ models are also
well-defined:

$\displaystyle\min_{\Phi}\ $
$\displaystyle\left\\{\mathbb{E}_{w}[R(w\circ\Phi)]+\lambda(\mathbb{E}_{w}[|\nabla_{w}R(w\circ\Phi)|])^{2}\right\\},$
(IRM-TV-$\ell_{1}$) $\displaystyle\min_{\Phi}\ $
$\displaystyle\\{\mathbb{E}_{w\leftarrow\frac{1_{(E)}}{E}}[R(w{\circ}\Phi)]{+}\lambda\max_{\rho}(\mathbb{E}_{w\leftarrow\rho}[|\nabla_{w}R(w{\circ}\Phi)|])^{2}\\}.$
(Minimax-TV-$\ell_{1}$)

The proof is presented in Appendix A.4. Note that
$(\mathbb{E}_{w}[|\nabla_{w}R(w\circ\Phi)|])^{2}=\|\nabla_{w}R(w\circ\Phi)\|_{1}^{2}$
if $\|\cdot\|_{1}$ is the $L^{1}$ norm for
$L^{1}(\Omega,\mathscr{F}_{\Omega},\nu)$, while
$\mathbb{E}_{w}[|\nabla_{w}R(w\circ\Phi)|^{2}]=\|\nabla_{w}R(w\circ\Phi)\|_{2}^{2}$
if $\|\cdot\|_{2}$ is the $L^{2}$ norm for
$L^{2}(\Omega,\mathscr{F}_{\Omega},\nu)$. Hence we retain a squared
TV-$\ell_{1}$ term in (IRM-TV-$\ell_{1}$) to be consistent with the squared
TV-$\ell_{2}$ term in (IRM-TV-$\ell_{2}$), so as (Minimax-TV-$\ell_{1}$).

The term $|\nabla_{w}R(w\circ\Phi)|$ in (IRM-TV-$\ell_{1}$) and (Minimax-
TV-$\ell_{1}$) is non-differentiable w.r.t. $\Phi$. Therefore, conventional
backpropagation or other gradient-type methods cannot be used in this model.
We develop a closed-form subgradient computation to solve (IRM-TV-$\ell_{1}$)
and (Minimax-TV-$\ell_{1}$), shown in Appendix B.1. It operates similarly to a
gradient computation and will not increase computational complexity.

The following proposition ensures that IRM-TV-$\ell_{1}$ has the coarea
formula, which produces a learning risk that is blocky (piecewise-constant)
w.r.t. the environment. It explains the mathematical essence of why IRM-
TV-$\ell_{1}$ is able to learn invariant features.

######  Proposition 3.6.

Assume the induced measure $\nu$ for $w$ to be the Lebesgue measure on
$\mathbb{R}^{d}$. Given Conditions 1$\sim$3 and 4b, the following coarea
formula holds:

$\displaystyle\int_{\Omega}|\nabla_{w}R(w\circ\Phi)|\,\mathrm{d}\nu=\int_{-\infty}^{\infty}\int_{\\{w\in\Omega:R(w\circ\Phi)=\gamma\\}}\,\mathrm{d}s\,\mathrm{d}\gamma,$
(10)

where $s$ is the $(d-1)$-dimensional Hausdorff measure (i.e., area in $(d-1)$
dimensions).

The proof is given in Appendix A.5.

###  3.3 Out-of-distribution Generalization

OOD generalization (Ben-Tal et al., 2009) refers to the ability of a trained
model to be generalized to an unseen domain. It can be formally defined as
follows (Arjovsky et al., 2020; Sagawa et al., 2020; Duchi & Namkoong, 2021;
Krueger et al., 2021; Lin et al., 2022):

$\displaystyle\min_{\Phi}\max_{e\in\mathcal{E}_{all}}R(\Phi,e),$ (OOD)

where $\mathcal{E}_{all}$ denotes the global environment set that contains all
the possible environments that can occur in the test, and $\Phi$ denotes the
model parameters for $R$. Two facts are implied in this framework: 1.
$\max_{e\in\mathcal{E}_{all}}R(\Phi,e)<\infty$ for some $\Phi$, or else it
cannot be minimized by any $\Phi$. 2. Both $\mathop{\rm
arg\,max}_{e\in\mathcal{E}_{all}}R(\Phi,e)$, $\forall\Phi$ and $\mathop{\rm
arg\,min}_{\Phi}\max_{e\in\mathcal{E}_{all}}R(\Phi,e)$ exist, or else the
worst-case environment or the optimal model parameters cannot be identified.
To absorb the environment variable $e$ into the classifier $w$, Condition 2
should be strengthened:

Condition 2a. The representation of the risk metric $R$ by $e$ only lies on
the classifier $w$, i.e., $R(\Phi,e)\equiv R(w(e)\circ\Phi)$.

We assume that the following Theorems 3.7, 3.8 and 3.9 satisfy Conditions
1,2a,3 and 4b on the probability space
$(\mathcal{E}_{all},\mathscr{F}_{all},\mu_{all})$.

######  Theorem 3.7.

(OOD) is equivalent to

$\displaystyle\min_{\Phi}\max_{w\in\Omega_{all}}R(w\circ\Phi),$ (OOD-$w$)

where $\Omega_{all}=w(\mathcal{E}_{all})$.

The proof is given in Appendix A.6. In fact, IRM, ZIN and TIVA implicitly
exploit this theorem to set up their models.

On the other hand, (IRM-TV-$\ell_{1}$) under the global environment set
becomes

$\displaystyle\min_{\Phi}\ $
$\displaystyle\left\\{\mathbb{E}_{w}^{all}[R(w\circ\Phi)]+\lambda(\mathbb{E}_{w}^{all}[|\nabla_{w}R(w\circ\Phi)|])^{2}\right\\},$
(IRM-TV-$\ell_{1}$-global)

where $\mathbb{E}_{w}^{all}$ denotes the expectation w.r.t. $w$ on the induced
probability space $(\Omega_{all},\mathscr{F}_{\Omega_{all}},\nu_{all})$. Our
first step to establish OOD generalization is to investigate the conditions
under which (IRM-TV-$\ell_{1}$-global) is really equivalent to (OOD-$w$).

######  Theorem 3.8 (IRM-TV-$\ell_{1}$-global Achieving OOD Generalization).

1) The penalty parameter $\lambda$ should be allowed to vary with $\Phi$.
Otherwise, (IRM-TV-$\ell_{1}$-global) cannot achieve (OOD-$w$).

2) For each $\Phi$, if $\mathbb{E}_{w}^{all}[|\nabla_{w}R(w{\circ}\Phi)|]>0$,
then there exists some $\lambda_{\Phi}\geqslant 0$ depending only on $\Phi$
such that

$\displaystyle\max_{w\in\Omega_{all}}R(w{\circ}\Phi)$
$\displaystyle{=}\mathbb{E}_{w}^{all}[R(w{\circ}\Phi)]{+}\lambda_{\Phi}(\mathbb{E}_{w}^{all}[|\nabla_{w}R(w{\circ}\Phi)|])^{2}.$
(11)

If $\mathbb{E}_{w}^{all}[|\nabla_{w}R(w{\circ}\Phi)|]=0$, (3.8) still holds
when $R(w{\circ}\Phi)$ is Lipschitz continuous w.r.t. $w$.

3) An optimal point $\Phi^{\bullet}$ of the following model

$\displaystyle\min_{\Phi}\ $
$\displaystyle\left\\{\mathbb{E}_{w}^{all}[R(w\circ\Phi)]+\lambda_{\Phi}(\mathbb{E}_{w}^{all}[|\nabla_{w}R(w\circ\Phi)|])^{2}\right\\}$
(IRM-TV-$\ell_{1}$-global-$\Phi$)

is also an optimal point of (OOD-$w$), and vice versa.

The proof is given in Appendix A.7. Item 1) allows $\lambda$ to vary with
$\Phi$. Without this variation, even a simple function fitting task cannot be
generalized, as shown in Appendix A.7. Item 2) ensures the existence of
$\lambda_{\Phi}$ that fills the gap between the maximum and the expectation of
$R(w\circ\Phi)$. Item 3) establishes the equivalence between the optimal
points of (IRM-TV-$\ell_{1}$-global) and (OOD-$w$).

######  Theorem 3.9 (Minimax-TV-$\ell_{1}$-global Achieving OOD
Generalization).

1) The penalty parameter $\lambda$ should be allowed to vary with $\Phi$.
Otherwise, the following model cannot achieve (OOD-$w$):

$\displaystyle\min_{\Phi}\ $
$\displaystyle\\{\mathbb{E}_{w\leftarrow\frac{1_{(E)}}{E}}^{all}[R(w{\circ}\Phi)]{+}\lambda\max_{\rho}(\mathbb{E}_{w\leftarrow\rho}^{all}[|\nabla_{w}R(w{\circ}\Phi)|])^{2}\\}.$
(Minimax-TV-$\ell_{1}$-global)

2) For each $\Phi$, if
$\mathbb{E}_{w\leftarrow\rho}^{all}[|\nabla_{w}R(w{\circ}\Phi)|]>0$ for some
$\rho$, then there exists some $\lambda_{\Phi}\geqslant 0$ depending only on
$\Phi$ such that

$\displaystyle\max_{w\in\Omega_{all}}R(w{\circ}\Phi)$
$\displaystyle{=}\mathbb{E}_{w\leftarrow\frac{1_{(E)}}{E}}^{all}[R(w{\circ}\Phi)]{+}\lambda_{\Phi}\max_{\rho}(\mathbb{E}_{w\leftarrow\rho}^{all}[|\nabla_{w}R(w{\circ}\Phi)|])^{2}.$
(12)

If $\mathbb{E}_{w\leftarrow\rho}^{all}[|\nabla_{w}R(w{\circ}\Phi)|]=0$ for all
$\rho$, (3.8) still holds when $R(w{\circ}\Phi)$ is Lipschitz continuous
w.r.t. $w$.

3) An optimal point $\Phi^{\bullet}$ of the following model

$\displaystyle\min_{\Phi}$
$\displaystyle\\{\mathbb{E}_{w\leftarrow\frac{1_{(E)}}{E}}^{all}[R(w{\circ}\Phi)]{+}\lambda_{\Phi}\max_{\rho}(\mathbb{E}_{w\leftarrow\rho}^{all}[|\nabla_{w}R(w{\circ}\Phi)|])^{2}\\}$
(Minimax-TV-$\ell_{1}$-global-$\Phi$)

is also an optimal point of (OOD-$w$), and vice versa.

The proof is given in Appendix A.8.

Theorems 3.7, 3.8, and 3.9 reveal the conditions under which TV-$\ell_{1}$
models can achieve OOD generalization under the global environment set. We
then investigate whether the training environment set can be generalized to
the global environment set. In most cases, we encounter some fundamental
environments in the training set that form a basis for environment
representation, either explicitly (e.g., IRM, REx) or implicitly (e.g., ZIN,
TIVA).

######  Definition 3.10 (Basis from Training Environments).

$\displaystyle\mathcal{B}_{tr}:=\\{e_{i}\\}_{i\in E},$ (13)

where $E$ denotes the index set for this basis. $E$ can be finite, countable
or uncountable.

Normally, $\mathcal{B}_{tr}\subseteq\mathcal{E}_{all}$, and usually
$\mathcal{B}_{tr}\subsetneqq\mathcal{E}_{all}$ because there may be unseen
environments outside the training set. This is why OOD generalization methods
try to expand $\mathcal{B}_{tr}$ from different approaches. An intuitive and
popular approach is to consider the linear space spanned by
$\mathcal{B}_{tr}$:

$\displaystyle\mathcal{E}_{tr}^{\mathscr{L}}:=\mathscr{L}(\mathcal{B}_{tr}).$
(14)

REx (Krueger et al., 2021) explicitly takes this form for the training
environment set (see Appendix A.2). Although ZIN and TIVA aim to expand
$\mathcal{B}_{tr}$ to its convex hull, they actually take
$\mathcal{E}_{tr}^{\mathscr{L}}$ as well and adopt a probability measure
$\rho$ on it according to Theorem 3.4.

The following theorem reveals the minimum requirements for global environment
set generalization.

######  Theorem 3.11.

Under the same conditions as in Theorem 3.5,
$\mathcal{E}_{all}\in\mathscr{F}_{tr}$ is the minimum requirement for (IRM-
TV-$\ell_{1}$) and (Minimax-TV-$\ell_{1}$) to be generalized to (IRM-
TV-$\ell_{1}$-global-$\Phi$) and (Minimax-TV-$\ell_{1}$-global-$\Phi$),
respectively.

The proof is given in Appendix A.9. This theorem reveals two important facts:

1\. $\mathcal{E}_{tr}$ should be abundant and diverse enough. In particular,
$\mathcal{E}_{tr}\supseteq\mathcal{E}_{all}$ is necessary. It may seem
impossible at first thought, but one can try to expand $\mathcal{B}_{tr}$ as
demonstrated above, like $\mathcal{E}_{tr}^{\mathscr{L}}$. After all, a model
can hardly foresee the environment of the cosmos in $\mathcal{E}_{all}$ with
only grass and deserts in $\mathcal{B}_{tr}$. But it may be able to learn the
environment with half deserts and half grass, or even the environment of a
forest (see Figure 1). The maximum $\sigma$-algebra of $\mathcal{E}_{tr}$ is
its power set $2^{\mathcal{E}_{tr}}$. If $\mathcal{E}_{all}\notin
2^{\mathcal{E}_{tr}}$, then distortion is inevitable in (IRM-
TV-$\ell_{1}$-global-$\Phi$) and (Minimax-TV-$\ell_{1}$-global-$\Phi$).
Recently, the spurious feature diversification strategy (Lin et al., 2023) has
been proposed to expand $\mathcal{E}_{tr}$, which is consistent with the above
analysis.

Figure 1: A simple example for the linear environment space
$\mathcal{E}_{tr}^{\mathscr{L}}$ (the dashed line).

2\. The measure $\mu$ for $e$ on $(\mathcal{E}_{tr},\mathscr{F}_{tr})$ should
be accurate enough. Even if $\mathcal{E}_{tr}\supseteq\mathcal{E}_{all}$
cannot guarantee a measurable $\mathcal{E}_{all}$. For example, let
$\mathcal{E}_{tr}\supsetneqq\mathcal{E}_{all}\supsetneqq\emptyset$ and the
measure $\mu$ be the coarsest one such that
$\mathscr{F}_{tr}=\\{\emptyset,\mathcal{E}_{tr}\\}$, then
$\mu(\mathcal{E}_{all})$ does not exist. For another example, let
$\mathcal{E}_{tr}=\mathbb{R}^{E}$, then the most accurate measure that can be
realized by modern computing technologies is the Lebesgue measure $\mu_{L}$
defined on the Lebesgue $\sigma$-algebra $\mathscr{F}_{L}$. However,
$\mathscr{F}_{L}\subsetneqq 2^{\mathbb{R}^{E}}$ and there are non-measurable
sets such as the Vitali set. Hence if
$\mathcal{E}_{all}\notin\mathscr{F}_{L}$, distortion is also inevitable in
(IRM-TV-$\ell_{1}$-global-$\Phi$) and (Minimax-TV-$\ell_{1}$-global-$\Phi$).
It makes little sense to only expand $\mathcal{E}_{tr}$ without equipping it
with an accurate $\mu$.

Theorem 3.11 actually deploys $\mathscr{F}_{tr}$ as the $\sigma$-algebra for
$\mathcal{E}_{all}$. However, $\mathcal{E}_{all}$ may have its own information
structure, represented by its own $\sigma$-algebra $\mathscr{F}_{all}$ and
measure $\mu_{all}$. Then it will be more complicated to explore (OOD-$w$),
leading to the following sufficient requirements.

######  Corollary 3.12.

Under the same conditions as in Theorem 3.5,
$\mathscr{F}_{tr}\supseteq\mathscr{F}_{all}$ is the sufficient requirement for
(IRM-TV-$\ell_{1}$) and (Minimax-TV-$\ell_{1}$) to be generalized to (IRM-
TV-$\ell_{1}$-global-$\Phi$) and (Minimax-TV-$\ell_{1}$-global-$\Phi$),
respectively.

To summarize this subsection, there are three ways to improve OOD
generalization: 1. Allow $\lambda_{\Phi}$ to vary with $\Phi$. 2. Expand the
training environment basis $\mathcal{B}_{tr}$ to a larger training environment
set $\mathcal{E}_{tr}$ by all means. 3. Equip $\mathcal{E}_{tr}$ with a more
accurate measure $\mu$.

##  4 Experiments

We assess the proposed IRM-TV-$\ell_{1}$ framework by comparisons with
existing OOD methods in simulation and real-world experiments, following (Lin
et al., 2022; Tan et al., 2023). Seven state-of-the-art methods: IRM (Arjovsky
et al., 2020), groupDRO (Sagawa et al., 2020), EIIL (Creager et al., 2021),
LfF (Nam et al., 2020), HRM (Liu et al., 2021), ZIN (Lin et al., 2022), and
TIVA (Tan et al., 2023), as well as the baseline Empirical risk minimization
(ERM) are taken into comparisons. IRM, groupDRO, and IRM-TV-$\ell_{1}$ require
ground-truth environment partitions, while Minimax-TV-$\ell_{1}$ is the
counterpart of IRM-TV-$\ell_{1}$ for non-environment-partition situations.
Implementation details are presented in Appendix B. Each experiment is
repeated for 10 times to compute the averaged result and the standard
deviation (STD) for each method.

###  4.1 Simulation Study

Each synthetic data set is characterized by temporal heterogeneity with
distributional shift w.r.t time and has been used for OOD generalization
evaluation in (Lin et al., 2022; Tan et al., 2023). For time $t\in[0,1]$, the
interested binary outcome is denoted by $Y(t)$, whose cause and effect are the
invariant features $X_{v}(t)\in\mathbb{R}$ and the spurious features
$X_{s}(t)\in\mathbb{R}$, respectively (see Appendix B.2 for more details). The
correlation between $X_{v}(t)$ and $Y(t)$ is stable and controlled by a
parameter $p_{v}$, while the correlation between $X_{s}(t)$ and $Y(t)$ is
changing over $t$ and controlled by a parameter $p_{s}(t)$. In training sample
generation, $p_{s}(t)$ is fixed as $p_{s}^{-}$ for $t\in[0,0.5)$ and as
$p_{s}^{+}$ for $t\in[0.5,1]$. Thus, the data generation process follows the
settings of triplet $(p_{s}^{-},p_{s}^{+},p_{v})$. Evaluation on four test
environments is performed with $p_{s}\in\\{0.999,0.8,0.2,0.001\\}$ and $p_{v}$
unchanged. Time $t$ is used as the auxiliary variable in ZIN and Minimax-
TV-$\ell_{1}$ for environment inference.

We present the mean accuracy and the worst accuracy over the four test
environments in Table 1. When environment partitions are unavailable, Minimax-
TV-$\ell_{1}$ outperforms the other competitors in all but one case where
$(p_{s}^{-},p_{s}^{+})=(0.999,0.7)$ and $p_{v}(t)=0.8$. When environment
partitions are available, IRM-TV-$\ell_{1}$ is the best method. As the
correlation between $X_{s}(t)$ and $Y(t)$ grows, the TV-$\ell_{1}$-based
models are more robust in two aspects: 1. The gap between the mean and the
worst performance of the TV-$\ell_{1}$-based models is generally smaller than
that of the other competitors; 2. The performance of the TV-$\ell_{1}$-based
models is less affected by such change.

We also show how the TV-$\ell_{1}$-based models improve identifiability of
invariant features over the TV-$\ell_{2}$-based models in Appendix C.

Table 1: Upper: accuracies (%) of competing methods on four test environments in simulation study averaged by 10 repetitions. Lower: STDs computed by 10 repetitions. Env partition |  $(p_{s}^{-},p_{s}^{+})$ | (0.999, 0.7) | (0.999, 0.8) | (0.999, 0.9)
---|---|---|---|---
$p_{v}(t)$ | 0.9 | 0.8 | 0.9 | 0.8 | 0.9 | 0.8
Method |  Mean |  Worst |  Mean |  Worst |  Mean |  Worst |  Mean |  Worst |  Mean |  Worst |  Mean |  Worst
False |  ERM |  76.22 |  58.81 |  59.80 |  25.95 |  69.34 |  43.06 |  55.96 |  15.60 |  60.62 |  23.30 |  53.10 |  8.04
EIIL |  39.43 |  18.22 |  64.95 |  48.45 |  50.26 |  47.02 |  68.86 |  54.91 |  61.33 |  52.70 |  69.82 |  58.58
HRM |  76.52 |  59.78 |  59.98 |  26.97 |  69.87 |  44.49 |  56.40 |  16.85 |  60.57 |  23.46 |  53.16 |  8.37
TIVA |  82.54 |  76.74 |  75.82 |  70.97 |  81.53 |  73.05 |  69.78 |  56.23 |  71.42 |  49.95 |  59.47 |  30.77
ZIN |  87.70 |  85.86 |  78.33 |  76.60 |  86.78 |  84.86 |  77.42 |  75.12 |  83.42 |  78.62 |  74.03 |  67.45
Minmax-TV-$\ell_{1}$ |  88.67 |  87.83 |  78.14 |  76.68 |  88.55 |  87.62 |  78.74 |  77.56 |  87.01 |  85.74 |  77.31 |  74.54
True |  groupDRO |  72.42 |  54.90 |  63.74 |  43.37 |  71.09 |  51.60 |  62.78 |  40.21 |  69.67 |  47.72 |  61.81 |  36.44
IRM |  87.84 |  86.20 |  78.33 |  76.58 |  86.84 |  84.42 |  77.48 |  74.80 |  84.16 |  77.89 |  74.53 |  68.72
IRM-TV-$\ell_{1}$ |  88.03 |  86.40 |  78.49 |  76.88 |  87.10 |  84.90 |  77.95 |  75.65 |  84.84 |  80.06 |  75.55 |  70.77
False |  ERM |  1.17 |  2.06 |  1.04 |  2.06 |  1.23 |  2.47 |  0.76 |  1.42 |  1.10 |  2.01 |  0.62 |  0.95
EIIL |  1.52 |  3.18 |  1.46 |  1.72 |  1.70 |  3.09 |  1.43 |  2.26 |  2.46 |  1.99 |  1.58 |  2.04
HRM |  1.35 |  2.71 |  0.94 |  2.43 |  0.75 |  1.83 |  0.71 |  2.33 |  0.84 |  1.29 |  0.45 |  0.93
TIVA |  6.12 |  11.09 |  3.55 |  7.18 |  4.83 |  9.19 |  6.46 |  13.96 |  5.18 |  10.34 |  6.32 |  13.66
ZIN |  1.05 |  2.19 |  1 |  1.43 |  1.67 |  2.73 |  1.43 |  2.13 |  3.52 |  6.72 |  2.09 |  3.86
Minmax-TV-$\ell_{1}$ |  0.57 |  0.60 |  0.84 |  1.03 |  0.45 |  0.50 |  0.67 |  0.74 |  1.28 |  1.66 |  0.65 |  1.13
True |  groupDRO |  8.45 |  18.08 |  6.99 |  16.84 |  8.42 |  19.03 |  6.71 |  17.27 |  8.27 |  18.51 |  6.52 |  16.45
IRM |  0.82 |  2.01 |  0.91 |  1.49 |  1.16 |  2.34 |  1.82 |  3.01 |  1.98 |  4.11 |  3.14 |  4.52
IRM-TV-$\ell_{1}$ |  0.86 |  2.08 |  0.74 |  1.33 |  1.35 |  2.67 |  1.24 |  2.22 |  2.19 |  4.77 |  2.92 |  4.31

###  4.2 Real-world Experiments

House Price Prediction. We use the House Prices data
set111https://www.kaggle.com/c/house-prices-advanced-regression-
techniques/data to verify the TV-$\ell_{1}$-based models in a regression task.
In this experiment, $15$ variables including the number of bathrooms,
locations, etc., are used to predict the house price. Samples with built year
in period $[1900,1950]$ are used for training and those with built year in
period $(1950,2000]$ are used for test. The house price is normalized within
the same built year. The built year is used as an auxiliary variable in both
ZIN and Minimax-TV-$\ell_{1}$ for environment inference. The training samples
are divided into $5$ segments with $10$-year range in each segment. Then each
segment is considered as having the same environment.

Table 2 reports the mean squared errors (MSE) of competing methods in this
regression task. IRM-TV-$\ell_{1}$ achieves the best results in the
environment-partition case, and Minimax-TV-$\ell_{1}$ achieves the best
results in the non-environment-partition case. Again, the TV-$\ell_{1}$-based
models have smaller gaps between the mean and the worst performance than the
other competitors.

Table 2: Upper: mean squared errors of competing methods in house price
prediction averaged by 10 repetitions. Lower: STDs computed by 10 repetitions.
Env Partition |  Method |  Train |  Test |  Worst
---|---|---|---|---
False |  ERM |  0.1057 |  0.4409 |  0.6206
EIIL |  0.1103 |  0.3939 |  0.5581
HRM |  0.5578 |  0.5949 |  0.7250
TIVA |  0.2575 |  0.4418 |  0.6145
ZIN |  0.2241 |  0.4293 |  0.6198
Minmax-TV-$\ell_{1}$ |  0.2168 |  0.3395 |  0.4983
True |  groupDRO |  0.1271 |  0.7358 |  1.0611
IRM |  0.5663 |  0.8168 |  1.1168
IRM-TV-$\ell_{1}$ |  0.3261 |  0.4420 |  0.6096
False |  ERM |  0.0017 |  0.0435 |  0.0641
EIIL |  0.0020 |  0.0305 |  0.0460
HRM |  0.0593 |  0.0025 |  0.0052
TIVA |  0.0002 |  0.0019 |  0.0062
ZIN |  0.1137 |  0.1994 |  0.2869
Minmax-TV-$\ell_{1}$ |  0.0652 |  0.0638 |  0.0958
True |  groupDRO |  0.0029 |  0.0877 |  0.1287
IRM |  0.1389 |  0.3115 |  0.4511
IRM-TV-$\ell_{1}$ |  0.1279 |  0.2503 |  0.3342

CelebA. This data set contains face images of celebrities (Liu et al., 2015).
The task is to identify the smiling faces, which is deliberately correlated
with gender. $512$-dimensional deep features of face images are extracted by a
pre-trained ResNet18 (He et al., 2016), and the invariant features are learned
by subsequent multilayer perceptrons. Seven additional descriptive variables
including _Young_ , _Blond Hair_ , _Eyeglasses_ , _High Cheekbones_ , _Big
Nose_ , _Bags Under Eyes_ , and _Chubby_ are fed into ZIN and Minimax-
TV-$\ell_{1}$ for environment inference. The gender variable is only used as
the environment indicator for groupDRO, IRM, and IRM-TV-$\ell_{1}$.

Table 3 presents the results. The TV-$\ell_{1}$-based models achieve the best
accuracies in the mean and the worst scenarios whether environment partitions
are available or not. Moreover, Minimax-TV-$\ell_{1}$ is closer to IRM and
IRM-TV-$\ell_{1}$ than other environment inference methods (EIIL, LfF, ZIN,
and TIVA). It indicates that TV-$\ell_{1}$ narrows the gap between absence and
presence of environment information.

Table 3: Upper: accuracies (%) of competing methods on CelebA averaged by 10
repetitions. Lower: STDs computed by 10 repetitions.
Env Partition |  Method |  Train |  Test |  Worst
---|---|---|---|---
False |  ERM |  63.76 |  63.99 |  62.05
EIIL |  59.12 |  58.15 |  54.22
LfF |  57.50 |  57.73 |  56.18
TIVA |  64.36 |  64.23 |  61.63
ZIN |  78.32 |  76.73 |  76.19
Minmax-TV-$\ell_{1}$ |  85.12 |  83.68 |  81.45
True |  groupDRO |  81.50 |  81.19 |  79.27
IRM |  85.59 |  82.54 |  80.75
IRM-TV-$\ell_{1}$ |  84.79 |  83.47 |  81.21
False |  ERM |  14.45 |  14.16 |  14.16
EIIL |  8.74 |  8.48 |  10.23
LfF |  0.12 |  0.24 |  0.57
TIVA |  1.68 |  1.99 |  1.47
ZIN |  1.16 |  0.87 |  0.85
Minmax-TV-$\ell_{1}$ |  0.92 |  0.33 |  0.43
True |  groupDRO |  0.31 |  0.48 |  0.74
IRM |  1.49 |  1.35 |  0.99
IRM-TV-$\ell_{1}$ |  0.59 |  0.48 |  0.67

Landcover. The Landcover data set records time series and the corresponding
land cover types from the satellite data (Gislason et al., 2006; Russwurm et
al., 2020; Xie et al., 2021). Time series data with dimension $46\times 8$ are
used as input to identify one from six land cover types. The invariant feature
extractor $\Phi$ is instantiated as a 1D-CNN to handle the time series input,
following (Xie et al., 2021; Lin et al., 2022). Ground-truth environment
partitions are unavailable, thus latitude and longitude are taken as auxiliary
information for environment inference. All methods are trained on non-African
data, and then tested on both non-African (not overlapping the training data)
and African data. Corresponding results are denoted by IID Test and OOD Test.

As shown in Table 4, Minimax-TV-$\ell_{1}$ achieves the best performance in
all of the IID, OOD and Worst Tests, with at least $1.4\%$ higher than the
second best competitor. Hence Minimax-TV-$\ell_{1}$ can be better generalized
to unseen environments.

Table 4: Upper: accuracies (%) of competing methods on Landcover averaged by
10 repetitions. Lower: STDs computed by 10 repetitions.

Method | Train | IID Test | OOD Test | Worst
---|---|---|---|---
ERM | 66.61 | 66.44 | 61.54 | 60.80
EIIL | 64.11 | 63.81 | 60.43 | 59.53
LfF | 58.12 | 57.89 | 55.76 | 55.07
TIVA | 67.49 | 64.79 | 52.02 | 51.46
ZIN | 70.02 | 69.42 | 62.22 | 61.87
Minmax-TV-$\ell_{1}$ | 73.59 | 71.95 | 63.77 | 63.25
ERM | 1.82 | 1.56 | 0.92 | 0.77
EIIL | 1.66 | 1.72 | 0.88 | 1.21
LfF | 2.73 | 2.45 | 1.96 | 1.93
TIVA | 0.28 | 0.62 | 0.98 | 1.09
ZIN | 1.09 | 1.14 | 1.09 | 1.21
Minmax-TV-$\ell_{1}$ | 0.69 | 0.63 | 1.17 | 1.37

Adult Income Prediction. In this task we use the Adult data
set222https://archive.ics.uci.edu/dataset/2/adult to predict if the income of
an individual exceeds $50K/yr based on the census data. We split the data set
into four subgroups regarded as separated environments according to
$race\in\\{\text{Black},\text{Non-Black}\\}$ and
$sex\in\\{\text{Male},\text{Female}\\}$. We randomly choose two thirds of data
from the subgroups Black Male and Non-Black Female for training, and then
verify models across all four subgroups with the rest data. Six integral
variables: _Age_ , _FNLWGT_ , _Eduction-Number_ , _Capital-Gain_ , _Capital-
Loss_ , and _Hours-Per-Week_ are fead into ZIN and Minimax-TV-$\ell_{1}$ for
environment inference. Ground-truth environment indicators are provided for
groupDRO, IRM and IRM-TV-$\ell_{1}$. Categorical variables except race and sex
are encoded by one-hot coding, followed by the principal component analysis
transform with over 99% cumulative explained variance ratio kept. The
transformed features are combined with the integral variables, yielding
59-dimensional representations, which are subsequently normalized to have zero
mean and unit variance for invariant feature learning.

Results are shown in Table 5. Minimax-TV-$\ell_{1}$ and IRM-TV-$\ell_{1}$
achieve the best accuracies in the mean and the worst scenarios within their
respective categories. Again, Minimax-TV-$\ell_{1}$ is closer to IRM and IRM-
TV-$\ell_{1}$ than other environment inference methods. Hence Minimax-
TV-$\ell_{1}$ narrows the gap between absence and presence of environment
information.

Table 5: Upper: accuracies (%) of competing methods in adult income prediction averaged by 10 repetitions. Lower: STDs computed by 10 repetitions.  Env Partition |  Method |  Train |  Test |  Worst
---|---|---|---|---
False |  ERM |  93.34 |  82.16 |  79.55
EIIL |  79.97 |  72.77 |  70.94
LfF |  82.03 |  75.04 |  73.00
TIVA |  91.45 |  81.95 |  79.28
ZIN |  93.16 |  82.26 |  79.67
Minmax-TV-$\ell_{1}$ |  92.40 |  83.33 |  80.95
True |  groupDRO |  87.51 |  76.42 |  73.07
IRM |  93.19 |  82.32 |  79.76
IRM-TV-$\ell_{1}$ |  92.42 |  83.31 |  80.93
False |  ERM |  0.31 |  0.33 |  0.37
EIIL |  0.56 |  0.61 |  0.73
LfF |  5.54 |  3.01 |  2.45
TIVA |  0.12 |  0.39 |  0.46
ZIN |  0.17 |  0.27 |  0.29
Minmax-TV-$\ell_{1}$ |  0.11 |  0.14 |  0.16
True |  groupDRO |  0.59 |  1.29 |  1.54
IRM |  0.28 |  0.23 |  0.29
IRM-TV-$\ell_{1}$ |  0.19 |  0.18 |  0.19

##  5 Conclusion and Discussion

We theoretically show that IRM is essentially an IRM-TV-$\ell_{2}$ model that
simultaneously minimizes the total empirical risk and its total variation with
respect to the classifier variable. Following this idea, we propose the IRM-
TV-$\ell_{1}$ and the Minimax-TV-$\ell_{1}$ models for learning tasks with or
without environment partitions, respectively. They allow broader classes of
functions to be the learning risk and the feature extractor, and preserve
invariant features based on the coarea formula. Moreover, we investigate the
requirements for the TV-$\ell_{1}$ framework to achieve out-of-distribution
(OOD) generalization. Extensive experiments on both synthetic and real-world
data sets show that the TV-$\ell_{1}$ framework performs well in several OOD
tests.

A direct impact of this work is to provide a new technical approach that
identifies, analyzes and constructs different kinds of invariants with
nonsmooth or even discontinuous modules. Although these properties are
difficult to handle, they precisely show a reasonable nature of generalization
ability and robustness. Further improvements on this work may lie in designing
an adaptive penalty parameter to improve OOD generalization, constructing a
diverse and representative training environment space, and developing new
TV-$\ell_{1}$ models for deep learning. These are nontrivial tasks and we
shall put them in future works.

## Acknowledgments

We would like to thank the anonymous reviewers for their constructive
suggestions that help to improve this paper.

This work is supported in part by the National Natural Science Foundation of
China under grant 62176103, in part by the Science and Technology Planning
Project of Guangzhou under grants 2024A04J9896, 2024A04J4225, and in part by
the Fundamental Research Funds for the Central Universities under grant
21623341.

Code is available at https://github.com/laizhr/IRM-TV.

## Impact Statement

This paper presents work whose goal is to advance the field of Machine
Learning. There are no potential societal consequences of our work that we
feel must be specifically highlighted here.

## References

  * Arjovsky et al. (2020) Arjovsky, M., Bottou, L., Gulrajani, I., and Lopez-Paz, D.  Invariant risk minimization, 2020.
  * Ben-Tal et al. (2009) Ben-Tal, A., Ghaoui, L. E., and Nemirovski, A.  _Robust Optimization_ , volume 28.  Princeton University Press, 2009.
  * Chan & Esedoglu (2005) Chan, T. and Esedoglu, S.  Aspects of total variation regularized $l^{1}$ function approximation.  _SIAM Journal on Applied Mathematics_ , (65):1817–1837, 2005.
  * Chan et al. (2006) Chan, T., Esedoglu, S., Park, F., and Yip, A.  _Total Variation Image Restoration: Overview and Recent Developments_ , pp. 17–31.  Springer US, Boston, MA, 2006.
  * Chen et al. (2006) Chen, T., Yin, W., Zhou, X. S., Comaniciu, D., and Huang, T. S.  Total variation models for variable lighting face recognition.  _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , 28(9):1519–1524, 2006.
  * Creager et al. (2021) Creager, E., Jacobsen, J.-H., and Zemel, R.  Environment inference for invariant learning.  In Meila, M. and Zhang, T. (eds.), _Proceedings of the 38th International Conference on Machine Learning_ , volume 139, pp. 2189–2200, 18–24 Jul 2021.
  * Duchi & Namkoong (2021) Duchi, J. C. and Namkoong, H.  Learning models with uniform performance via distributionally robust optimization.  _The Annals of Statistics_ , 49(3):1378 – 1406, 2021.
  * Federer (1959) Federer, H.  Curvature measures.  _Transactions of the American Mathematical Society_ , 93(3):418–491, 1959.
  * Gislason et al. (2006) Gislason, P. O., Benediktsson, J. A., and Sveinsson, J. R.  Random forests for land cover classification.  _Pattern Recognition Letters_ , 27(4):294–300, 2006.
  * He et al. (2016) He, K., Zhang, X., Ren, S., and Sun, J.  Deep residual learning for image recognition.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , June 2016.
  * Kingma & Ba (2015) Kingma, D. P. and Ba, J.  Adam: A method for stochastic optimization.  In _International Conference on Learning Representations_ , 2015.
  * Krueger et al. (2021) Krueger, D., Caballero, E., Jacobsen, J.-H., Zhang, A., Binas, J., Zhang, D., Priol, R. L., and Courville, A.  Out-of-distribution generalization via risk extrapolation (rex).  In _Proceedings of the 38th International Conference on Machine Learning_ , pp. 5815–5826, 2021.
  * Lin et al. (2022) Lin, Y., Zhu, S., Tan, L., and Cui, P.  Zin: When and how to learn invariance without environment partition?  In _Advances in Neural Information Processing Systems_ , pp. 24529–24542, 2022.
  * Lin et al. (2023) Lin, Y., Tan, L., Hao, Y., Wong, H., Dong, H., Zhang, W., Yang, Y., and Zhang, T.  Spurious feature diversification improves out-of-distribution generalization, 2023.
  * Liu et al. (2021) Liu, J., Hu, Z., Cui, P., Li, B., and Shen, Z.  Heterogeneous risk minimization.  In _Proceedings of the 38th International Conference on Machine Learning_ , volume 139, pp. 6804–6814, 18–24 Jul 2021.
  * Liu et al. (2015) Liu, Z., Luo, P., Wang, X., and Tang, X.  Deep learning face attributes in the wild.  In _Proceedings of the IEEE International Conference on Computer Vision (ICCV)_ , December 2015.
  * Mumford & Shah (1985) Mumford, D. and Shah, J.  Boundary detection by minimizing functionals, I.  In _Proceedings of IEEE Conference on Computer Vision and Pattern Recognition_ , pp. 22–26, 1985.
  * Nam et al. (2020) Nam, J., Cha, H., Ahn, S., Lee, J., and Shin, J.  Learning from failure: De-biasing classifier from biased classifier.  In _Advances in Neural Information Processing Systems_ , volume 33, pp. 20673–20684, 2020.
  * Recht et al. (2019) Recht, B., Roelofs, R., Schmidt, L., and Shankar, V.  Do ImageNet classifiers generalize to ImageNet?  In _Proceedings of the 36th International Conference on Machine Learning_ , pp. 5389–5400, 2019.
  * Rudin et al. (1992) Rudin, L. I., Osher, S., and Fatemi, E.  Nonlinear total variation based noise removal algorithms.  _Physica D: Nonlinear Phenomena_ , 60(1):259–268, 1992.
  * Russwurm et al. (2020) Russwurm, M., Wang, S., Korner, M., and Lobell, D.  Meta-learning for few-shot land cover classification.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) Workshops_ , June 2020.
  * Sagawa et al. (2020) Sagawa, S., Koh, P. W., Hashimoto, T. B., and Liang, P.  Distributionally robust neural networks.  In _8th International Conference on Learning Representations, ICLR_ , 2020.
  * Tan et al. (2023) Tan, X., Yong, L., Zhu, S., Qu, C., Qiu, X., Yinghui, X., Cui, P., and Qi, Y.  Provably invariant learning without domain information.  In _Proceedings of the 40th International Conference on Machine Learning_ , pp. 33563–33580, 2023.
  * Xie et al. (2021) Xie, S. M., Kumar, A., Jones, R., Khani, F., Ma, T., and Liang, P.  In-n-out: Pre-training and self-training using auxiliary information for out-of-distribution robustness.  In _International Conference on Learning Representations_ , 2021.
  * Zhou et al. (2022) Zhou, X., Lin, Y., Zhang, W., and Zhang, T.  Sparse invariant risk minimization.  In _Proceedings of the 39th International Conference on Machine Learning_ , pp. 27222–27244, 2022.

##  Appendix A Proofs

###  A.1 Proof of Theorem 3.1

####  A.1.1 Induced Measure $\nu$

First, we verify that there is a measure $\nu$ for $w$ induced by $\mu$. From
Condition 3, let $\Omega=w(\mathcal{E}_{tr})$ and
$\mathscr{F}_{\Omega}:=\\{w(\mathcal{F}):\mathcal{F}\in\mathscr{F}_{tr}\\}$.
By definition, $\mathscr{F}_{\Omega}$ is the $\sigma$-algebra for $\Omega$
under mapping $w$. Define a set function $\nu$ as follows:

$\displaystyle\nu:\mathscr{F}_{\Omega}\rightarrow\mathbb{R},\quad\nu(\mathcal{G}):=\mu(w^{-1}(\mathcal{G})).$
(15)

It is well defined because $w^{-1}(\mathcal{G})\in\mathscr{F}_{tr}$. Now we
prove that $\nu$ is a measure. Since $\mu$ is a measure, the following non-
negativity and zero measure are obvious:

$\displaystyle\forall\
\mathcal{G}\in\mathscr{F}_{\Omega},\quad\nu(\mathcal{G})=\mu(w^{-1}(\mathcal{G}))\geqslant
0,\quad\nu(\emptyset)=\mu(w^{-1}(\emptyset))=\mu(\emptyset)=0.$ (16)

To prove countable additivity, let $\\{\mathcal{G}_{i}\\}_{i=1}^{\infty}$ be
any countable collection of mutually disjoint sets in $\mathscr{F}_{\Omega}$.
From the property of measurable function,
$\\{w^{-1}(\mathcal{G}_{i})\\}_{i=1}^{\infty}$ is also a countable collection
of mutually disjoint sets in $\mathscr{F}_{tr}$. Therefore,

$\displaystyle\nu(\bigcup_{i=1}^{\infty}\mathcal{G}_{i})=\mu(w^{-1}(\bigcup_{i=1}^{\infty}\mathcal{G}_{i}))=\mu(\bigcup_{i=1}^{\infty}[w^{-1}(\mathcal{G}_{i})])=\sum_{i=1}^{\infty}\mu(w^{-1}(\mathcal{G}_{i}))=\sum_{i=1}^{\infty}\nu(\mathcal{G}_{i}),$
(17)

where the mutually disjoint $\\{\mathcal{G}_{i}\\}_{i=1}^{\infty}$ yields the
second equality and the mutually disjoint
$\\{w^{-1}(\mathcal{G}_{i})\\}_{i=1}^{\infty}$ yields the third equality,
respectively. Combining (16) and (17), we can conclude that $\nu$ is a
measure.

####  A.1.2 Constructing Measurable Functions $R(w\circ\Phi)$ and
$|\nabla_{w}R(w\circ\Phi)|$

From Conditions 2 and 4,
$R(w(e)\circ\Phi):\mathcal{E}_{tr}\rightarrow\mathbb{R}$ is a measurable
function on $(\mathcal{E}_{tr},\mathscr{F}_{tr})$ and

$\displaystyle\int_{\mathcal{E}_{tr}}R(w\circ\Phi,e)\,\mathrm{d}\mu=\int_{\mathcal{E}_{tr}}R(w(e)\circ\Phi)\,\mathrm{d}\mu\in(-\infty,+\infty).$
(18)

We can directly construct a function
$R(w\circ\Phi):\Omega\rightarrow\mathbb{R}$ on $(\Omega,\mathscr{F}_{\Omega})$
and verify its measurability:

$\displaystyle\forall\omega\in\Omega,\quad
R(\omega\circ\Phi):=R(w[w^{-1}(\omega)]\circ\Phi),$ (19)

where $w^{-1}(\omega)\subseteq\mathcal{E}_{tr}$. From Condition 3, we have
$\mathcal{E}=w^{-1}(\mathcal{G})\in\mathscr{F}_{tr}$ for any
$\mathcal{G}\in\mathscr{F}_{\Omega}$. Hence

$\displaystyle
R(\mathcal{G}\circ\Phi):\mathscr{F}_{\Omega}\rightarrow\mathbb{R},\quad
R(\mathcal{G}\circ\Phi):=R(w[w^{-1}(\mathcal{G})]\circ\Phi)=R(w(\mathcal{E})\circ\Phi)$
(20)

is well-defined. Next, let $\mathcal{R}_{0}\in\mathscr{F}_{\mathbb{R}}$, where
$\mathscr{F}_{\mathbb{R}}$ denotes the $\sigma$-algebra of all the Lebesgue
measurable sets contained in $\mathbb{R}$. Then there exists some
$\mathcal{E}_{0}\in\mathscr{F}_{tr}$ such that

$\displaystyle R(w(\mathcal{E}_{0})\circ\Phi)=\mathcal{R}_{0}.$ (21)

Let $\mathcal{G}_{0}:=w(\mathcal{E}_{0})\in\mathscr{F}_{\Omega}$. Since
$\mathcal{R}_{0}$ is arbitrary, we have

$\displaystyle\mathring{\mathscr{F}}_{\Omega}:=\\{\mathcal{G}_{0}:\
R(\mathcal{G}_{0}\circ\Phi)=\mathcal{R}_{0},\
\mathcal{R}_{0}\in\mathscr{F}_{\mathbb{R}}\\}\subseteq\mathscr{F}_{\Omega}.$
(22)

It remains to verify that $\mathring{\mathscr{F}}_{\Omega}$ is a
$\sigma$-algebra, forming a sub-$\sigma$-algebra of $\mathscr{F}_{\Omega}$. It
is evident that $\emptyset,\Omega\in\mathring{\mathscr{F}}_{\Omega}$. For any
$\mathcal{G}_{0}\in\mathring{\mathscr{F}}_{\Omega}$,

$\displaystyle\Omega\backslash\mathcal{G}_{0}=w[w^{-1}(\Omega\backslash\mathcal{G}_{0})]=w[w^{-1}(\Omega)\backslash
w^{-1}(\mathcal{G}_{0})]=w[\mathcal{E}_{tr}\backslash
w^{-1}(\mathcal{G}_{0})].$ (23)

Because
$\mathcal{G}_{0}\in\mathring{\mathscr{F}}_{\Omega}\subseteq\mathscr{F}_{\Omega}$,
$w^{-1}(\mathcal{G}_{0})\in\mathscr{F}_{tr}$. Since $\mathscr{F}_{tr}$ is
closed under complement, $\mathcal{E}_{tr}\backslash
w^{-1}(\mathcal{G}_{0})\in\mathscr{F}_{tr}$. Thus
$w[\mathcal{E}_{tr}\backslash
w^{-1}(\mathcal{G}_{0})]\in{\mathscr{F}}_{\Omega}$. Then

$\displaystyle
R((\Omega\backslash\mathcal{G}_{0})\circ\Phi)=R(w[\mathcal{E}_{tr}\backslash
w^{-1}(\mathcal{G}_{0})]\circ\Phi)=\tilde{\mathcal{R}}_{0}\ \text{for some}\
\tilde{\mathcal{R}}_{0}\in\mathscr{F}_{\mathbb{R}}.$ (24)

The first equality follows from (23) and the second equality follows from the
measurability of $R(w(e)\circ\Phi)$ w.r.t. $e$. (24) indicates that
$\Omega\backslash\mathcal{G}_{0}\in\mathring{\mathscr{F}}_{\Omega}$ according
to (22). Hence $\mathring{\mathscr{F}}_{\Omega}$ is closed under complement.
Finally, let any
$\\{\mathcal{G}_{i}\in\mathring{\mathscr{F}}_{\Omega}\\}_{i=1}^{\infty}$. Then

$\displaystyle\bigcup_{i=1}^{\infty}\mathcal{G}_{i}=w\left[w^{-1}(\bigcup_{i=1}^{\infty}\mathcal{G}_{i})\right]=w\left[\bigcup_{i=1}^{\infty}w^{-1}(\mathcal{G}_{i})\right].$
(25)

Because
$\mathcal{G}_{i}\in\mathring{\mathscr{F}}_{\Omega}\subseteq\mathscr{F}_{\Omega}$,
$w^{-1}(\mathcal{G}_{i})\in\mathscr{F}_{tr}$. Since $\mathscr{F}_{tr}$ is
closed under countable unions,
$\bigcup_{i=1}^{\infty}w^{-1}(\mathcal{G}_{i})\in\mathscr{F}_{tr}$. Thus
$w\left[\bigcup_{i=1}^{\infty}w^{-1}(\mathcal{G}_{i})\right]\in{\mathscr{F}}_{\Omega}$.
Then

$\displaystyle
R\left(\left(\bigcup_{i=1}^{\infty}\mathcal{G}_{i}\right)\circ\Phi\right)=R\left(w\left[\bigcup_{i=1}^{\infty}w^{-1}(\mathcal{G}_{i})\right]\circ\Phi\right)=\mathcal{R}_{\infty}\
\text{for some}\ \mathcal{R}_{\infty}\in\mathscr{F}_{\mathbb{R}}.$ (26)

The first equality follows from (25) and the second equality follows again
from the measurability of $R(w(e)\circ\Phi)$ w.r.t. $e$. (26) indicates that
$\bigcup_{i=1}^{\infty}\mathcal{G}_{i}\in\mathring{\mathscr{F}}_{\Omega}$
according to (22). Hence $\mathring{\mathscr{F}}_{\Omega}$ is closed under
countable unions. To summarize, $\mathring{\mathscr{F}}_{\Omega}$ satisfies
the definition of $\sigma$-algebra and forms a sub-$\sigma$-algebra of
$\mathscr{F}_{\Omega}$.

On the other hand, the measurable function $|\nabla_{w}R(w\circ\Phi)|$ can be
constructed via the same procedure. Condition 2 implies that
$\nabla_{w}R(w\circ\Phi,e)\,\mathrm{d}\mu\equiv\nabla_{w}R(w(e)\circ\Phi)\,\mathrm{d}\mu$
since the partial differential is taken on $w$. Moreover, Condition 4
indicates that

$\displaystyle\int_{\mathcal{E}_{tr}}|\nabla_{w}R(w\circ\Phi,e)|^{2}\,\mathrm{d}\mu=\int_{\mathcal{E}_{tr}}|\nabla_{w}R(w(e)\circ\Phi)|^{2}\,\mathrm{d}\mu\in(-\infty,+\infty).$
(27)

We are now ready to establish the integrations w.r.t. $w$ from (18) and (27).

####  A.1.3 Integrations w.r.t. $w$

The integrations w.r.t. $w$ can be established via the integrations w.r.t.
$e$. First, we aim to establish the following:

$\displaystyle\int_{\Omega}R(w\circ\Phi)\,\mathrm{d}\nu:=\int_{\mathcal{E}_{tr}}R(w(e)\circ\Phi)\,\mathrm{d}\mu,\quad\forall\
R(w(e)\circ\Phi)\in L^{1}(\mathcal{E}_{tr},\mathscr{F}_{tr},\mu).$ (28)

A.1.1 and A.1.2 have already verified the eligible measure $\,\mathrm{d}\nu$
and the measurable function $R(w\circ\Phi)$. We start from the simplest case.
Let $\mathcal{G}$ be any measurable set in $\mathscr{F}_{\Omega}$ and
$\mathcal{E}=w^{-1}(\mathcal{G})\in\mathscr{F}_{tr}$. Define the indicator
function of $\mathcal{G}$ as follows:

$\displaystyle\mathbb{I}_{\mathcal{G}}(w)=\left\\{\begin{array}[]{ll}1&\text{if}\quad
w\in\mathcal{G},\\\ 0&\text{if}\quad w\notin\mathcal{G}.\end{array}\right.$
(31)

We have

$\displaystyle\int_{\Omega}\mathbb{I}_{\mathcal{G}}(w)\,\mathrm{d}\nu=\nu(\mathcal{G})=\mu(\mathcal{E})=\int_{\mathcal{E}_{tr}}\mathbb{I}_{\mathcal{E}}(e)\,\mathrm{d}\mu.$
(32)

This relationship also holds for $a\mathbb{I}_{\mathcal{G}}(w)$, $\forall
a\geqslant 0$. Next, define the sets of non-negative simple functions as

$\displaystyle\mathcal{S}^{+}_{\Omega}:=\left\\{\sum_{i=1}^{m}a_{i}\mathbb{I}_{\mathcal{G}_{i}}(w):m\in\mathbb{N}^{+},a_{i}\geqslant
0,\mathcal{G}_{i}\in\mathscr{F}_{\Omega}\right\\}\quad\text{and}\quad\mathcal{S}^{+}_{tr}:=\left\\{\sum_{i=1}^{m}a_{i}\mathbb{I}_{\mathcal{E}_{i}}(e):m\in\mathbb{N}^{+},a_{i}\geqslant
0,\mathcal{E}_{i}\in\mathscr{F}_{tr}\right\\}.$ (33)

Then for any $\sum_{i=1}^{m}a_{i}\mathbb{I}_{\mathcal{G}_{i}}(w)$, denote
$\mathcal{E}_{i}=w^{-1}(\mathcal{G}_{i})$, $i=1,\cdots,m$. We have

$\displaystyle\int_{\Omega}\sum_{i=1}^{m}a_{i}\mathbb{I}_{\mathcal{G}_{i}}(w)\,\mathrm{d}\nu=\sum_{i=1}^{m}a_{i}\int_{\Omega}\mathbb{I}_{\mathcal{G}_{i}}(w)\,\mathrm{d}\nu=\sum_{i=1}^{m}a_{i}\nu(\mathcal{G}_{i})=\sum_{i=1}^{m}a_{i}\mu(\mathcal{E}_{i})=\sum_{i=1}^{m}a_{i}\int_{\mathcal{E}_{tr}}\mathbb{I}_{\mathcal{E}_{i}}(e)\,\mathrm{d}\mu=\int_{\mathcal{E}_{tr}}\sum_{i=1}^{m}a_{i}\mathbb{I}_{\mathcal{E}_{i}}(e)\,\mathrm{d}\mu.$
(34)

That is, for any
$\sum_{i=1}^{m}a_{i}\mathbb{I}_{\mathcal{G}_{i}}(w)\in\mathcal{S}^{+}_{\Omega}$,
there exists some
$\sum_{i=1}^{m}a_{i}\mathbb{I}_{\mathcal{E}_{i}}(e)\in\mathcal{S}^{+}_{tr}$
satisfying (34).

Next, suppose $R^{+}(w\circ\Phi)$ is a non-negative measurable function w.r.t.
$w$. Then it can be approached by a sequence of functions in
$\mathcal{S}^{+}_{\Omega}$. Specifically, let

$\displaystyle a_{i}:=\frac{i}{2^{n}},\mathcal{G}_{i}:=\\{w:a_{i}\leqslant
R(w\circ\Phi)<a_{i+1}\\},\ i=0,1,\cdots,n2^{n}-1;\ \mathcal{G}_{R\geqslant
n}:=\\{w:R(w\circ\Phi)\geqslant n\\}.$ (35)

Define the sequence of approaching functions as follows:

$\displaystyle
R^{+}_{n}(w\circ\Phi)=\sum_{i=0}^{n2^{n}-1}a_{i}\mathbb{I}_{\mathcal{G}_{i}}(w)+n\mathbb{I}_{\mathcal{G}_{R\geqslant
n}}(w),\quad n=1,2,3,\cdots.$ (36)

It can be easily observed that
$R^{+}_{n}(w\circ\Phi)\in\mathcal{S}^{+}_{\Omega}$ and
$R^{+}_{n}(w\circ\Phi)\uparrow R^{+}(w\circ\Phi)$ (point-wise monotonically
non-decreasing and convergent). Let $\mathcal{E}_{i}=w^{-1}(\mathcal{G}_{i})$,
$i=0,1,\cdots,n2^{n}-1$, and $\mathcal{E}_{R\geqslant
n}=w^{-1}(\mathcal{G}_{R\geqslant n})$. Then

$\displaystyle
R^{+}_{n}(w(e)\circ\Phi)=\sum_{i=0}^{n2^{n}-1}a_{i}\mathbb{I}_{\mathcal{E}_{i}}(e)+n\mathbb{I}_{\mathcal{E}_{R\geqslant
n}}(e),\quad n=1,2,3,\cdots.$ (37)

We also have $R^{+}_{n}(w(e)\circ\Phi)\in\mathcal{S}^{+}_{tr}$ and
$R^{+}_{n}(w(e)\circ\Phi)\uparrow R^{+}(w(e)\circ\Phi)$ for some non-negative
measurable function $R^{+}(w(e)\circ\Phi)$ w.r.t. $e$. From (34),

$\displaystyle\int_{\Omega}R^{+}_{n}(w\circ\Phi)\,\mathrm{d}\nu=\int_{\mathcal{E}_{tr}}R^{+}_{n}(w(e)\circ\Phi)\,\mathrm{d}\mu,\quad
n=1,2,3,\cdots.$ (38)

Moreover, this sequence of integrals is monotonically increasing and thus
converges to
$\int_{\mathcal{E}_{tr}}R^{+}(w(e)\circ\Phi)\,\mathrm{d}\mu\in\bar{\mathbb{R}}$
(the extended real space including $+\infty$ and $-\infty$). We can directly
define

$\displaystyle\int_{\Omega}R^{+}(w\circ\Phi)\,\mathrm{d}\nu:=\int_{\mathcal{E}_{tr}}R^{+}(w(e)\circ\Phi)\,\mathrm{d}\mu.$
(39)

Since we only consider integrable functions $R^{+}(w(e)\circ\Phi)\in
L^{1}(\mathcal{E}_{tr},\mathscr{F}_{tr},\mu)$ by Condition 4, (39) indicates
that $\int_{\Omega}R^{+}(w\circ\Phi)\,\mathrm{d}\nu<+\infty$.

Last, suppose $R(w\circ\Phi)$ is a measurable function w.r.t. $w$. Following
conventional methodology, it can be decomposed into a positive part and a
negative part:

$\displaystyle R(w\circ\Phi)=R^{+}(w\circ\Phi)-R^{-}(w\circ\Phi),\quad
R^{+}(w\circ\Phi):=\max\\{R(w\circ\Phi),0\\},\
R^{-}(w\circ\Phi):=\max\\{-R(w\circ\Phi),0\\}.$ (40)

Both $R^{+}(w\circ\Phi)$ and $R^{-}(w\circ\Phi)$ are non-negative measurable
functions. By (39), there exist non-negative measurable functions
$R^{+}(w(e)\circ\Phi)$ and $R^{-}(w(e)\circ\Phi)$ w.r.t. $e$ such that

$\displaystyle\int_{\Omega}R^{+}(w\circ\Phi)\,\mathrm{d}\nu:=\int_{\mathcal{E}_{tr}}R^{+}(w(e)\circ\Phi)\,\mathrm{d}\mu,\quad\int_{\Omega}R^{-}(w\circ\Phi)\,\mathrm{d}\nu:=\int_{\mathcal{E}_{tr}}R^{-}(w(e)\circ\Phi)\,\mathrm{d}\mu.$
(41)

Let $R(w(e)\circ\Phi)=R^{+}(w(e)\circ\Phi)-R^{-}(w(e)\circ\Phi)$. Again, we
only consider $R(w(e)\circ\Phi)\in
L^{1}(\mathcal{E}_{tr},\mathscr{F}_{tr},\mu)$ by Condition 4, thus both
integrals in (41) are finite. Define

$\displaystyle\int_{\Omega}R(w\circ\Phi)\,\mathrm{d}\nu:=\int_{\Omega}R^{+}(w\circ\Phi)\,\mathrm{d}\nu-\int_{\Omega}R^{-}(w\circ\Phi)\,\mathrm{d}\nu,$
(42)

then $\int_{\Omega}R(w\circ\Phi)\,\mathrm{d}\nu\in(-\infty,+\infty)$.

One can apply the above procedure to $|\nabla_{w}R(w\circ\Phi)|^{2}$ and
exploit $|\nabla_{w}R(w\circ\Phi)|\in
L^{2}(\mathcal{E}_{tr},\mathscr{F}_{tr},\mu)$ from Condition 4 to define

$\displaystyle\int_{\Omega}|\nabla_{w}R(w\circ\Phi)|^{2}\,\mathrm{d}\nu\in[0,+\infty).$
(43)

Note that $|\nabla_{w}R(w\circ\Phi)|$ is a non-negative measurable function.

####  A.1.4 Expectations w.r.t. $w$ and Well-defined (IRM-TV-$\ell_{2}$)

To ensure that $\int_{\Omega}R(w\circ\Phi)\,\mathrm{d}\nu$ and
$\int_{\Omega}|\nabla_{w}R(w\circ\Phi)|^{2}\,\mathrm{d}\nu$ are mathematical
expectations, we need to verify that $\nu$ is also a probability measure
induced by $\mu$:

$\displaystyle\nu(\Omega)=\mu(\mathcal{E}_{tr})=1,$ (44)

where $\mu(\mathcal{E}_{tr})=1$ because $\mu$ is a probability measure.
$\mathbb{E}_{w}[R(w\circ\Phi)]$ and
$\mathbb{E}_{w}[|\nabla_{w}R(w\circ\Phi)|^{2}]$ denote the two integrals when
$\nu$ is a probability measure. Since both integrals are finite,
$\mathbb{E}_{w}[R(w\circ\Phi)+\lambda|\nabla_{w}R(w\circ\Phi)|^{2}]$ is also
finite and can be minimized w.r.t. $\Phi$. Hence (IRM-TV-$\ell_{2}$) is well-
defined.

For an extension, one can establish a probability measure as long as
$\nu(\Omega)\in(0,+\infty)$:

$\displaystyle\tilde{\nu}(\cdot):=\nu(\cdot)/\nu(\Omega).$ (45)

It is very useful when $\Omega$ is a bounded subset of $\mathbb{R}^{d}$ and
$\nu$ is the Lebesgue measure.

###  A.2 Proof of Proposition 3.3

Given Conditions 1 and 4a, $\mathbb{E}_{e}[R(\Phi,e)]$ and
$\mathbb{E}_{e}[R^{2}(\Phi,e)]$ are well-defined and finite. We can use the
integral with probability measure $\mu$ to represent the first term in
(V-REx):

$\displaystyle\sum_{e\in\mathcal{E}_{tr}}R(\Phi,e):=\mathbb{E}_{e}[R(\Phi,e)]\in(-\infty,+\infty).$
(46)

Additionally, the variance of $R(\Phi,e)$ in (V-REx) can be generalized to

$\displaystyle\mathbb{V}(\\{R(\Phi,e)\\}_{e\in\mathcal{E}_{tr}}):=\mathbb{E}_{e}[(R(\Phi,e){-}\mathbb{E}_{e}[R(\Phi,e)])^{2}]=\mathbb{E}_{e}[R^{2}(\Phi,e)]-(\mathbb{E}_{e}[R(\Phi,e)])^{2}\in(-\infty,+\infty).$
(47)

Therefore, (V-REx-$\ell_{2}$) is well-defined and finite, and can be
minimized.

To illustrate how V-REx-$\ell_{2}$ extrapolates risk to a larger region, we
need to examine the graph space of the learning risk:
$\mathcal{E}_{tr}\times\mathbb{R}:=\\{(e,R(\Phi,e)):e\in\mathcal{E}_{tr}\\}$.
$\mathbb{E}_{e}[R(\Phi,e)]$ can be seen as the mean risk of all the
environments, which is located in the center. Under some mild assumptions
(such as the continuity of $R(\Phi,e)$ w.r.t. $e$),
$\mathbb{E}_{e}[R(\Phi,e)]$ falls into the range of $R(\Phi,e)$. Then there
exists some $\tilde{e}$ such that
$R(\Phi,\tilde{e})=\mathbb{E}_{e}[R(\Phi,e)]$. This $\tilde{e}$ can be seen as
some mean-risk environment, and $(\tilde{e},\mathbb{E}_{e}[R(\Phi,e)])$ is in
the midst of the graph $\\{(e,R(\Phi,e)):e\in\mathcal{E}_{tr}\\}$. The
variance in (47) actually involves a set of vectors in the graph space
$\\{(e,R(\Phi,e))-(\tilde{e},\mathbb{E}_{e}[R(\Phi,e)]):e\in\mathcal{E}_{tr}\\}$
which can span a linear space to extrapolate risk, as shown in Figure A1. The
performance of such risk extrapolation depends on the diversity of the
training environment set $\mathcal{E}_{tr}$ and the measure $\mu$.

Figure A1: Risk extrapolation of V-REx-$\ell_{2}$ to a larger region.

###  A.3 Proof of Theorem 3.4

To prove this theorem, we first absorb (2.1) into the framework of this paper.
Suppose the learning model uses a finite set of $E$ fundamental environments
(not necessarily known by the model) as the basis defined in Definition 3.10:
$\mathcal{B}_{tr}:=\\{e_{i}\\}_{i=1}^{E}$. Then the first term in (2.1) can be
rewritten as follows:

$\displaystyle\tilde{R}(w{\circ}\Phi,\frac{1}{E}1_{(E)}){\cdot}1_{(E)}=\sum_{i=1}^{E}[\tilde{R}(w{\circ}\Phi,\frac{1}{E}1_{(E)})]_{i}=\frac{1}{E}\sum_{i=1}^{E}R(w{\circ}\Phi,e_{i}),$
(48)

where $[\cdot]_{i}$ denotes the $i$-th dimension of a vector. Here we use
$\tilde{R}$ to denote the original form of learning risk in (2.1), in order to
distinguish it from $R$ in our framework. In fact, $\tilde{R}$ is an
$E$-dimensional vector that represents the learning risk in the $E$
fundamental environments. This representation is equivalent to interpreting
$R$ as a function of $e_{i}$ with $e_{i}$ varying in the $E$ fundamental
environments. Therefore, (48) is true.

Similarly, the second term of (2.1) can be rewritten as follows:

$\displaystyle\|\nabla_{w}\tilde{R}(w{\circ}\Phi,\rho)\|_{2}^{2}=\sum_{i=1}^{E}[|\nabla_{w}\tilde{R}(w{\circ}\Phi,\rho)|]_{i}^{2}=\sum_{i=1}^{E}|\nabla_{w}R(w{\circ}\Phi,e_{i})|^{2}\rho_{i},\quad\rho\in\Delta^{E}.$
(49)

It further assigns a probability $\rho\in\Delta^{E}$ (see Eq. 2 for the
definition) to the fundamental environments, while (48) uses the uniform
probability $\frac{1}{E}1_{(E)}$.

Both (48) and (49) are essentially mathematical expectations:

$\displaystyle\mathbb{E}_{\frac{1_{(E)}}{E}}[R(w{\circ}\Phi,e)]:=\int_{\mathcal{E}_{tr}}R(w{\circ}\Phi,e)\,\mathrm{d}(\frac{1_{(E)}}{E})=\frac{1}{E}\sum_{i=1}^{E}R(w{\circ}\Phi,e_{i}),$
(50)
$\displaystyle\mathbb{E}_{\rho}[|\nabla_{w}R(w{\circ}\Phi,e)|^{2}]:=\int_{\mathcal{E}_{tr}}|\nabla_{w}R(w{\circ}\Phi,e)|^{2}\,\mathrm{d}\rho=\sum_{i=1}^{E}|\nabla_{w}R(w{\circ}\Phi,e_{i})|^{2}\rho_{i}.$
(51)

From Theorem 3.1, there exist induced probability measures for $w$ such that

$\displaystyle\mathbb{E}_{w\leftarrow\frac{1_{(E)}}{E}}[R(w{\circ}\Phi)]=\mathbb{E}_{\frac{1_{(E)}}{E}}[R(w{\circ}\Phi,e)],\quad\mathbb{E}_{w\leftarrow\rho}[|\nabla_{w}R(w{\circ}\Phi)|^{2}]=\mathbb{E}_{\rho}[|\nabla_{w}R(w{\circ}\Phi,e)|^{2}].$
(52)

Both ZIN and TIVA aim to learn the probability $\rho$ from the training
samples with environment information. They differ mainly in their learning
approach. Therefore, both ZIN and TIVA can be generalized to (Minimax-
TV-$\ell_{2}$).

###  A.4 Proof of Theorem 3.5

The proof is almost identical to that of A.1 and A.3, except that
$|\nabla_{w}R(w{\circ}\Phi)|^{2}$ is replaced by
$|\nabla_{w}R(w{\circ}\Phi)|$. Note that the condition
$|\nabla_{w}R(w(e)\circ\Phi)|\in L^{1}(\mathcal{E}_{tr},\mathscr{F}_{tr},\mu)$
in Theorem 3.5 is weaker than $|\nabla_{w}R(w(e)\circ\Phi)|\in
L^{2}(\mathcal{E}_{tr},\mathscr{F}_{tr},\mu)$ in Theorem 3.1 when
$\mu(\mathcal{E}_{tr})<\infty$. Therefore, (IRM-TV-$\ell_{1}$) and (Minimax-
TV-$\ell_{1}$) allow broader classes of $R(w{\circ}\Phi)$ to be used than
(IRM-TV-$\ell_{2}$) and (Minimax-TV-$\ell_{2}$).

###  A.5 Proof of Proposition 3.6

The key to prove this proposition is to partition the real space into
countable continuous intervals such that the level sets of $R(w\circ\Phi)$
w.r.t. these intervals have positive volumes. Specifically, we establish a set
of $K\in\mathbb{N}^{+}\cup\\{+\infty\\}$ intervals

$\displaystyle\\{B_{k}:=[a_{k},b_{k})\\}_{k=1}^{K}\quad\text{s.\ t.}\quad
a_{k}<b_{k}\leqslant a_{k+1},\quad k=1,2,\cdots,K,$ $\displaystyle\qquad
R(w\circ\Phi)\ \text{is Lipschitz continuous on}\
\mathcal{G}_{k}:=\\{w\in\Omega:a_{k}<R(w\circ\Phi)<b_{k}\\}.$ (53)

Since $a_{k}<b_{k}$, each $B_{k}$ contains at least one rational number.
Besides, since $B_{k}\cap B_{j}=\emptyset$ for any $k\neq j$, the rational
numbers in $B_{k}$ cannot overlap those in $B_{j}$. Since all the rational
numbers are dense in $\mathbb{R}$ and the volume $|B_{k}|:=(b_{k}-a_{k})>0$
for all $k$, there are at most countable intervals $\\{B_{k}\\}_{k=1}^{K}$. On
the contrary, if none of such interval $B_{k}$ exists, it is obvious that the
integral in (10) equals $0$.

Given the conditions,
$\int_{\Omega}|\nabla_{w}R(w\circ\Phi)|\,\mathrm{d}\nu<\infty$. Since
$B_{k}\subseteq\Omega$ and $|\nabla_{w}R(w\circ\Phi)|\geqslant 0$,
$\int_{B_{k}}|\nabla_{w}R(w\circ\Phi)|\,\mathrm{d}\nu\leqslant\int_{\Omega}|\nabla_{w}R(w\circ\Phi)|\,\mathrm{d}\nu<\infty$
holds for all $k$. Thus, we can decompose the integral as follows:

$\displaystyle\int_{\Omega}|\nabla_{w}R(w\circ\Phi)|\,\mathrm{d}\nu=\int_{\Omega\backslash(\bigcup_{k=1}^{K}\mathcal{G}_{k})}|\nabla_{w}R(w\circ\Phi)|\,\mathrm{d}\nu+\int_{\bigcup_{k=1}^{K}\mathcal{G}_{k}}|\nabla_{w}R(w\circ\Phi)|\,\mathrm{d}\nu.$
(54)

Since $\Omega\backslash(\bigcup_{k=1}^{K}\mathcal{G}_{k})$ corresponds to
$\mathbb{R}\backslash(\bigcup_{k=1}^{K}B_{k})$ where the function value
$R(w\circ\Phi)$ is discontinuous almost everywhere (a.e.), the Lebesgue
measure (i.e., the $d$-dimensional volume) $\,\mathrm{d}\nu=0$ on
$\Omega\backslash(\bigcup_{k=1}^{K}\mathcal{G}_{k})$. Hence, the first term in
(54)
$\int_{\Omega\backslash(\bigcup_{k=1}^{K}\mathcal{G}_{k})}|\nabla_{w}R(w\circ\Phi)|\,\mathrm{d}\nu=0$.
From (Federer, 1959), the coarea formula holds for all the Lipschitz
continuous parts:

$\displaystyle\int_{\mathcal{G}_{k}}|\nabla_{w}R(w\circ\Phi)|\,\mathrm{d}\nu=\int_{a_{k}}^{b_{k}}\int_{\\{w\in\Omega:R(w\circ\Phi)=\gamma\\}}\,\mathrm{d}s\,\mathrm{d}\gamma,\quad
k=1,2,\cdots,K.$ (55)

Exploiting the $\sigma$-additivity of Lebesgue measure yields

$\displaystyle\int_{\bigcup_{k=1}^{K}\mathcal{G}_{k}}|\nabla_{w}R(w\circ\Phi)|\,\mathrm{d}\nu$
$\displaystyle=\sum_{k=1}^{K}\int_{\mathcal{G}_{k}}|\nabla_{w}R(w\circ\Phi)|\,\mathrm{d}\nu=\sum_{k=1}^{K}\int_{a_{k}}^{b_{k}}\int_{\\{w\in\Omega:R(w\circ\Phi)=\gamma\\}}\,\mathrm{d}s\,\mathrm{d}\gamma$
$\displaystyle=\int_{\bigcup_{k=1}^{K}B_{k}}\int_{\\{w\in\Omega:R(w\circ\Phi)=\gamma\\}}\,\mathrm{d}s\,\mathrm{d}\gamma.$
(56)

Besides, it is also obvious to see that $\,\mathrm{d}\gamma=0$ on
$\mathbb{R}\backslash(\bigcup_{k=1}^{K}B_{k})$ where $\gamma=R(w\circ\Phi)$ is
discontinuous a.e. Hence
$\int_{\mathbb{R}\backslash(\bigcup_{k=1}^{K}B_{k})}\int_{\\{w\in\Omega:R(w\circ\Phi)=\gamma\\}}\,\mathrm{d}s\,\mathrm{d}\gamma=0$.
Together with (54) and (A.5), we have

$\displaystyle\int_{-\infty}^{\infty}\int_{\\{w\in\Omega:R(w\circ\Phi)=\gamma\\}}\,\mathrm{d}s\,\mathrm{d}\gamma$
$\displaystyle=\int_{\bigcup_{k=1}^{K}B_{k}}\int_{\\{w\in\Omega:R(w\circ\Phi)=\gamma\\}}\,\mathrm{d}s\,\mathrm{d}\gamma=\int_{\bigcup_{k=1}^{K}\mathcal{G}_{k}}|\nabla_{w}R(w\circ\Phi)|\,\mathrm{d}\nu$
$\displaystyle=\int_{\Omega}|\nabla_{w}R(w\circ\Phi)|\,\mathrm{d}\nu.$ (57)

It proves the coarea formula in Proposition 3.6.

This proof also explains the mathematical essence of why TV-$\ell_{1}$ models
can learn invariant features. To achieve a smaller TV for $R(w\circ\Phi)$, the
continuous interval $B_{k}$ should take a smaller length, especially on the
level set $\\{w\in\Omega:R(w\circ\Phi)=\gamma\\}$ with a large area. To do
this, the optimizing procedure may squash the graph of $R(w\circ\Phi)$ w.r.t.
$w$ in every continuous level set $\mathcal{G}_{k}$, making the graph more
blocky (piece-wise constant). This makes $R(w\circ\Phi)$ more robust w.r.t.
$w$ in each $\mathcal{G}_{k}$.

###  A.6 Proof of Theorem 3.7

From the implied facts of (OOD), we only need to prove that for any $\Phi$
such that $\max_{e\in\mathcal{E}_{all}}R(\Phi,e)<\infty$,
$\max_{e\in\mathcal{E}_{all}}R(\Phi,e)=\max_{w\in\Omega_{all}}R(w\circ\Phi)$.
From Conditions 2a and 3, it suffices to prove that $\forall\Phi$,
$\max_{e\in\mathcal{E}_{all}}R(w(e)\circ\Phi)=\max_{w\in\Omega_{all}}R(w\circ\Phi)$.
First, we examine how the left side yields the right side. Denote
$e^{*}=\mathop{\rm arg\,max}_{e\in\mathcal{E}_{all}}R(w(e)\circ\Phi)$. Then
$R(w(e^{*})\circ\Phi)\geqslant\max_{w\in\Omega_{all}}R(w\circ\Phi)$. If not,
then there exists $\omega^{\bullet}\in\Omega_{all}$ such that
$R(w(e^{*})\circ\Phi)<R(\omega^{\bullet}\circ\Phi)$. Since $w$ is a surjective
mapping, we have $w^{-1}(\omega^{\bullet})\neq\emptyset$ and $e^{*}\notin
w^{-1}(\omega^{\bullet})$. Picking up any $e^{\bullet}\in
w^{-1}(\omega^{\bullet})$, we have
$R(w(e^{\bullet})\circ\Phi)>R(w(e^{*})\circ\Phi)$, which violates the
optimality of $e^{*}$.

Second, we examine how the right side yields the left side. Denote
$w^{*}=\mathop{\rm arg\,max}_{w\in\Omega_{all}}R(w\circ\Phi)$. Then
$R(w^{*}\circ\Phi)\geqslant\max_{e\in\mathcal{E}_{all}}R(w(e)\circ\Phi)$. If
not, then there exists $e^{\bullet}\in\mathcal{E}_{all}$ such that
$R(w^{*}\circ\Phi)<R(w(e^{\bullet})\circ\Phi)$. Hence
$R(w^{*}\circ\Phi)<R(w^{\bullet}\circ\Phi)$ with
$w^{\bullet}:=w(e^{\bullet})\in\Omega_{all}$, which violates the optimality of
$w^{*}$.

Summarizing the inequalities of both directions, we have
$\max_{e\in\mathcal{E}_{all}}R(w(e)\circ\Phi)=\max_{w\in\Omega_{all}}R(w\circ\Phi)$,
$\forall\Phi$. Thus minimizing either side of this equality w.r.t. $\Phi$
leads to the same optimization model.

###  A.7 Proof of Theorem 3.8

####  A.7.1 Allowing $\lambda$ to Vary with $\Phi$

To prove that $\lambda$ should be allowed to vary with $\Phi$, we only need to
provide a counterexample with a fixed $\lambda$. Consider a simple function
fitting problem that uses $w\cdot\Phi$ to fit the constant function $-1$. We
use the absolute fitting error as the learning risk
$R(w\circ\Phi):=|w\cdot\Phi-(-1)|=|w\cdot\Phi+1|$. Suppose we aim to learn the
best feature parameter $\Phi\in[-1,1]$. The classifier $w$ lies uniformly in
$[-0.9,0.1]$, but its center deviates from $0$ after conveying the environment
information $e$. Equip $w$ with the uniform probability measure on
$[-0.9,0.1]$ as $\nu$. Now $R(w\circ\Phi)=1+w\cdot\Phi$ according to the
domains of $w$ and $\Phi$. It can be easily calculated that

$\displaystyle\max_{w\in\Omega_{all}}R(w{\circ}\Phi)=\max_{w\in[-0.9,0.1]}\\{1+w\cdot\Phi\\}=\left\\{\begin{array}[]{ll}1+0.1\Phi&\text{if}\quad\Phi\geqslant
0\\\ 1-0.9\Phi&\text{if}\quad\Phi<0\end{array}\right.,$ (60)
$\displaystyle\min_{\Phi\in[-1,1]}\max_{w\in\Omega_{all}}R(w{\circ}\Phi)=1,\quad\mathop{\rm
arg\,min}_{\Phi\in[-1,1]}\max_{w\in\Omega_{all}}R(w{\circ}\Phi)=0.$ (61)

However,

$\displaystyle\mathbb{E}_{w}^{all}[R(w\circ\Phi)]+\lambda(\mathbb{E}_{w}^{all}[|\nabla_{w}R(w\circ\Phi)|])^{2}=\int_{-0.9}^{0.1}(1+w\cdot\Phi)\,\mathrm{d}\nu+\lambda(|\Phi|\int_{-0.9}^{0.1}\,\mathrm{d}\nu)^{2}=1-\frac{2}{5}\Phi+\lambda\Phi^{2},$
(62)
$\displaystyle\min_{\Phi\in[-1,1]}\\{\mathbb{E}_{w}^{all}[R(w\circ\Phi)]+\lambda(\mathbb{E}_{w}^{all}[|\nabla_{w}R(w\circ\Phi)|])^{2}\\}=\left\\{\begin{array}[]{ll}1-\frac{1}{25\lambda}&\text{if}\quad\lambda>\frac{1}{5}\\\
\frac{3}{5}+\lambda&\text{if}\quad 0<\lambda\leqslant\frac{1}{5}\\\
\frac{3}{5}&\text{if}\quad\lambda=0\end{array}\right.,$ (66)
$\displaystyle\mathop{\rm
arg\,min}_{\Phi\in[-1,1]}\\{\mathbb{E}_{w}^{all}[R(w\circ\Phi)]+\lambda(\mathbb{E}_{w}^{all}[|\nabla_{w}R(w\circ\Phi)|])^{2}\\}=\left\\{\begin{array}[]{ll}\frac{1}{5\lambda}&\text{if}\quad\lambda>\frac{1}{5}\\\
1&\text{if}\quad 0\leqslant\lambda\leqslant\frac{1}{5}\end{array}\right..$
(69)

Comparing (66) with (61), we have

$\displaystyle\min_{\Phi\in[-1,1]}\\{\mathbb{E}_{w}^{all}[R(w\circ\Phi)]+\lambda(\mathbb{E}_{w}^{all}[|\nabla_{w}R(w\circ\Phi)|])^{2}\\}\neq\min_{\Phi\in[-1,1]}\max_{w\in\Omega_{all}}R(w{\circ}\Phi),$
$\displaystyle\mathop{\rm
arg\,min}_{\Phi\in[-1,1]}\\{\mathbb{E}_{w}^{all}[R(w\circ\Phi)]+\lambda(\mathbb{E}_{w}^{all}[|\nabla_{w}R(w\circ\Phi)|])^{2}\\}\neq\mathop{\rm
arg\,min}_{\Phi\in[-1,1]}\max_{w\in\Omega_{all}}R(w{\circ}\Phi),$ (70)

for any fixed $\lambda\geqslant 0$. Hence (IRM-TV-$\ell_{1}$-global) cannot
achieve (OOD-$w$) in terms of the objective value and the argument.

This counterexample can be explained as follows. The learning risk $R$
increases with $|\Phi|$ for the worst-case $w$. Therefore, $\Phi=0$ is the
optimal point in (OOD-$w$). However, the deviation of $w$ from $0$ also
deviates the expectation of $R$ w.r.t. $w$, leading to the deviation of
optimal $\Phi$ from $0$ in (IRM-TV-$\ell_{1}$-global). This deviation of
$\Phi$ cannot be offset by the TV term in (IRM-TV-$\ell_{1}$-global), no
matter how we set the fixed $\lambda$.

####  A.7.2 Existence of $\lambda_{\Phi}$

First, we consider the trivial case
$\mathbb{E}_{w}^{all}[|\nabla_{w}R(w\circ\Phi)|]=0$, which indicates a zero TV
term. When $R(w\circ\Phi)$ is Lipschitz continuous w.r.t. $w$,
$\mathbb{E}_{w}^{all}[|\nabla_{w}R(w\circ\Phi)|]=0$ implies that
$R(w\circ\Phi)$ is constant w.r.t. $w$. In this case,
$\mathbb{E}_{w}^{all}[R(w\circ\Phi)]=\max_{w\in\Omega_{all}}R(w{\circ}\Phi)$
and $\lambda$ can be set arbitrarily.

Next, we investigate the nontrivial case
$\mathbb{E}_{w}^{all}[|\nabla_{w}R(w\circ\Phi)|]>0$. We have

$\displaystyle\mathbb{E}_{w}^{all}[R(w\circ\Phi)]\leqslant\mathbb{E}_{w}^{all}[\max_{\omega\in\Omega_{all}}R(\omega{\circ}\Phi)]\leqslant(\max_{\omega\in\Omega_{all}}R(\omega{\circ}\Phi))\cdot\mathbb{E}_{w}^{all}[1]=\max_{\omega\in\Omega_{all}}R(\omega{\circ}\Phi).$
(71)

Note that $\max_{\omega\in\Omega_{all}}R(\omega{\circ}\Phi)$ is constant
w.r.t. $w$ and thus can be moved outside the expectation. Then
$\lambda_{\Phi}$ exists by the form:

$\displaystyle\lambda_{\Phi}:=\frac{(\max_{w\in\Omega_{all}}R(w{\circ}\Phi))-\mathbb{E}_{w}^{all}[R(w\circ\Phi)]}{(\mathbb{E}_{w}^{all}[|\nabla_{w}R(w\circ\Phi)|])^{2}}\geqslant
0.$ (72)

This quotient depends only on $\Phi$, but not $w$.

####  A.7.3 Achieving Optimality in (OOD-$w$)

$\Phi^{\bullet}$ being the optimal point of (IRM-TV-$\ell_{1}$-global-$\Phi$)
indicates

$\displaystyle\mathbb{E}_{w}^{all}[R(w\circ\Phi^{\bullet})]+\lambda_{\Phi}^{\bullet}(\mathbb{E}_{w}^{all}[|\nabla_{w}R(w\circ\Phi^{\bullet})|])^{2}\leqslant\mathbb{E}_{w}^{all}[R(w\circ\Phi)]+\lambda_{\Phi}(\mathbb{E}_{w}^{all}[|\nabla_{w}R(w\circ\Phi)|])^{2},\
\forall\Phi.$ (73)

In particular, $\Phi$ can be any optimal point $\Phi^{*}$ of (OOD-$w$).
Together with (3.8), we have

$\displaystyle\max_{w\in\Omega_{all}}R(w{\circ}\Phi^{\bullet})$
$\displaystyle=\mathbb{E}_{w}^{all}[R(w\circ\Phi^{\bullet})]+\lambda_{\Phi}^{\bullet}(\mathbb{E}_{w}^{all}[|\nabla_{w}R(w\circ\Phi^{\bullet})|])^{2}$
$\displaystyle\leqslant\mathbb{E}_{w}^{all}[R(w\circ\Phi^{*})]+\lambda_{\Phi}^{*}(\mathbb{E}_{w}^{all}[|\nabla_{w}R(w\circ\Phi^{*})|])^{2}=\max_{w\in\Omega_{all}}R(w{\circ}\Phi^{*}).$
(74)

Hence $\Phi^{\bullet}$ is also an optimal point of (OOD-$w$).

Conversely, $\Phi^{*}$ being the optimal point of (OOD-$w$) indicates

$\displaystyle\mathbb{E}_{w}^{all}[R(w\circ\Phi^{*})]$
$\displaystyle+\lambda_{\Phi}^{*}(\mathbb{E}_{w}^{all}[|\nabla_{w}R(w\circ\Phi^{*})|])^{2}=\max_{w\in\Omega_{all}}R(w{\circ}\Phi^{*})$
$\displaystyle\leqslant\max_{w\in\Omega_{all}}R(w{\circ}\Phi^{\bullet})=\mathbb{E}_{w}^{all}[R(w\circ\Phi^{\bullet})]+\lambda_{\Phi}^{\bullet}(\mathbb{E}_{w}^{all}[|\nabla_{w}R(w\circ\Phi^{\bullet})|])^{2}.$
(75)

Hence $\Phi^{*}$ is also an optimal point of (IRM-
TV-$\ell_{1}$-global-$\Phi$).

###  A.8 Proof of Theorem 3.9

####  A.8.1 Allowing $\lambda$ to Vary with $\Phi$

The counterexample is similar to that in A.7.1 except for some changes. First,
the measure $\nu_{\frac{1_{(E)}}{E}}$ for $w$ induced by $\frac{1_{(E)}}{E}$
can be set as the uniform probability measure on $[-0.9,0.1]$. The measure
$\nu_{\rho}$ for $w$ induced by $\rho$ can be any arbitrary probability
measure. Then

$\displaystyle\mathbb{E}_{w\leftarrow\frac{1_{(E)}}{E}}^{all}[R(w{\circ}\Phi)]{+}\lambda\max_{\rho}(\mathbb{E}_{w\leftarrow\rho}^{all}[|\nabla_{w}R(w{\circ}\Phi)|])^{2}$
$\displaystyle=$
$\displaystyle\int_{-0.9}^{0.1}(1+w\cdot\Phi)\,\mathrm{d}\nu_{\frac{1_{(E)}}{E}}+\lambda\max_{\rho}(|\Phi|\int_{-0.9}^{0.1}\,\mathrm{d}\nu_{\rho})^{2}=1-\frac{2}{5}\Phi+\lambda\Phi^{2},$
(76)

which is the same as that in (62). Similar results to (66) and (A.7.1) also
hold for any fixed $\lambda\geqslant 0$. Thus (OOD-$w$) cannot be achieved by
(Minimax-TV-$\ell_{1}$-global).

####  A.8.2 Existence of $\lambda_{\Phi}$

If $\mathbb{E}_{w\leftarrow\rho}^{all}[|\nabla_{w}R(w{\circ}\Phi)|]=0$ for all
$\rho$ and $R(w{\circ}\Phi)$ is Lipschitz continuous w.r.t. $w$, then
$R(w{\circ}\Phi)$ is constant w.r.t. $w$. In this case,
$\mathbb{E}_{w\leftarrow\frac{1_{(E)}}{E}}^{all}[R(w{\circ}\Phi)]=\max_{w\in\Omega_{all}}R(w{\circ}\Phi)$
and $\lambda$ can be set arbitrarily.

If $\mathbb{E}_{w\leftarrow\rho}^{all}[|\nabla_{w}R(w{\circ}\Phi)|]>0$ for
some $\rho$, then
$\max_{\rho}(\mathbb{E}_{w\leftarrow\rho}^{all}[|\nabla_{w}R(w{\circ}\Phi)|])^{2}>0$.
Similar to (71),

$\displaystyle\mathbb{E}_{w\leftarrow\frac{1_{(E)}}{E}}^{all}[R(w{\circ}\Phi)]\leqslant\max_{\omega\in\Omega_{all}}R(\omega{\circ}\Phi).$
(77)

Then $\lambda_{\Phi}$ exists by the form:

$\displaystyle\lambda_{\Phi}:=\frac{(\max_{w\in\Omega_{all}}R(w{\circ}\Phi))-\mathbb{E}_{w\leftarrow\frac{1_{(E)}}{E}}^{all}[R(w{\circ}\Phi)]}{\max_{\rho}(\mathbb{E}_{w\leftarrow\rho}^{all}[|\nabla_{w}R(w{\circ}\Phi)|])^{2}}\geqslant
0.$ (78)

This quotient depends only on $\Phi$, but not $w$.

####  A.8.3 Achieving Optimality in (OOD-$w$)

An optimal point $\Phi^{\bullet}$ of (Minimax-TV-$\ell_{1}$-global-$\Phi$)
satisfies

$\displaystyle\mathbb{E}_{w\leftarrow\frac{1_{(E)}}{E}}^{all}[R(w{\circ}\Phi^{\bullet})]{+}\lambda_{\Phi}^{\bullet}\max_{\rho}(\mathbb{E}_{w\leftarrow\rho}^{all}[|\nabla_{w}R(w{\circ}\Phi^{\bullet})|])^{2}\leqslant\mathbb{E}_{w\leftarrow\frac{1_{(E)}}{E}}^{all}[R(w{\circ}\Phi)]{+}\lambda_{\Phi}\max_{\rho}(\mathbb{E}_{w\leftarrow\rho}^{all}[|\nabla_{w}R(w{\circ}\Phi)|])^{2},\
\forall\Phi.$ (79)

$\Phi$ can be any optimal point $\Phi^{*}$ of (OOD-$w$). By combining (3.9),
we can deduce that

$\displaystyle\max_{w\in\Omega_{all}}R(w{\circ}\Phi^{\bullet})$
$\displaystyle=\mathbb{E}_{w\leftarrow\frac{1_{(E)}}{E}}^{all}[R(w{\circ}\Phi^{\bullet})]{+}\lambda_{\Phi}^{\bullet}\max_{\rho}(\mathbb{E}_{w\leftarrow\rho}^{all}[|\nabla_{w}R(w{\circ}\Phi^{\bullet})|])^{2}$
$\displaystyle\leqslant\mathbb{E}_{w\leftarrow\frac{1_{(E)}}{E}}^{all}[R(w{\circ}\Phi^{*})]{+}\lambda_{\Phi}^{*}\max_{\rho}(\mathbb{E}_{w\leftarrow\rho}^{all}[|\nabla_{w}R(w{\circ}\Phi^{*})|])^{2}=\max_{w\in\Omega_{all}}R(w{\circ}\Phi^{*}).$
(80)

Hence $\Phi^{\bullet}$ is also an optimal point of (OOD-$w$).

Conversely, an optimal point $\Phi^{*}$ of (OOD-$w$) satisfies

$\displaystyle\mathbb{E}_{w\leftarrow\frac{1_{(E)}}{E}}^{all}[R(w{\circ}\Phi^{*})]$
$\displaystyle+\lambda_{\Phi}^{*}\max_{\rho}(\mathbb{E}_{w\leftarrow\rho}^{all}[|\nabla_{w}R(w{\circ}\Phi^{*})|])^{2}=\max_{w\in\Omega_{all}}R(w{\circ}\Phi^{*})$
$\displaystyle\leqslant\max_{w\in\Omega_{all}}R(w{\circ}\Phi^{\bullet})=\mathbb{E}_{w\leftarrow\frac{1_{(E)}}{E}}^{all}[R(w{\circ}\Phi^{\bullet})]{+}\lambda_{\Phi}^{\bullet}\max_{\rho}(\mathbb{E}_{w\leftarrow\rho}^{all}[|\nabla_{w}R(w{\circ}\Phi^{\bullet})|])^{2}.$
(81)

Hence $\Phi^{*}$ is also an optimal point of (Minimax-
TV-$\ell_{1}$-global-$\Phi$).

###  A.9 Proof of Theorem 3.11

We denote the probability space for the training environment set as
$(\mathcal{E}_{tr},\mathscr{F}_{tr},\mu_{tr})$ and its induced probability
space for $w$ as $(\Omega_{tr},\mathscr{F}_{\Omega_{tr}},\nu_{tr})$. Since
$\mathcal{E}_{all}\in\mathscr{F}_{tr}$, we have
$\mathcal{E}_{all}\subseteq\mathcal{E}_{tr}$. Define $\mathscr{F}_{tr|all}$ as
the restricted $\sigma$-algebra on $\mathcal{E}_{all}$:

$\displaystyle\mathscr{F}_{tr|all}:=\\{A\cap\mathcal{E}_{all}:A\in\mathscr{F}_{tr}\\}.$
(82)

Since $\mathcal{E}_{all}\in\mathscr{F}_{tr}$, we have
$A\cap\mathcal{E}_{all}\in\mathscr{F}_{tr}$. Hence $\mathscr{F}_{tr|all}$ is a
sub-$\sigma$-algebra of $\mathscr{F}_{tr}$, and it can still use the
probability measure $\mu_{tr}$. Then
$(\mathcal{E}_{all},\mathscr{F}_{tr|all},\mu_{tr})$ can be used as a coarse
probability space for $\mathcal{E}_{all}$, and its induced probability space
for $w$ is $(\Omega_{all},\mathscr{F}_{\Omega_{tr|all}},\nu_{tr})$, where
$\Omega_{all}\subseteq\Omega_{tr}$ and
$\mathscr{F}_{\Omega_{tr|all}}\subseteq\mathscr{F}_{\Omega_{tr}}$.

Since the learning risk $R(w{\circ}\Phi)$ and $|\nabla_{w}R(w\circ\Phi)|$ are
defined on $(\Omega_{tr},\mathscr{F}_{\Omega_{tr}},\nu_{tr})$, the following
objective function and optimization model are well-defined:

$\displaystyle\int_{\Omega_{all}}R(w{\circ}\Phi)\,\mathrm{d}\nu_{tr}{+}\lambda_{\Phi}(\int_{\Omega_{all}}|\nabla_{w}R(w{\circ}\Phi)|\,\mathrm{d}\nu_{tr})^{2}=:\mathbb{E}_{w}^{tr}[R(w\circ\Phi)]+\lambda_{\Phi}(\mathbb{E}_{w}^{tr}[|\nabla_{w}R(w\circ\Phi)|])^{2},$
(83) $\displaystyle\min_{\Phi}\
\left\\{\mathbb{E}_{w}^{tr}[R(w\circ\Phi)]+\lambda_{\Phi}(\mathbb{E}_{w}^{tr}[|\nabla_{w}R(w\circ\Phi)|])^{2}\right\\}.$
(IRM-TV-$\ell_{1}$-training-$\Phi$)

The only difference of (IRM-TV-$\ell_{1}$-training-$\Phi$) from (IRM-
TV-$\ell_{1}$) is the variable $\lambda_{\Phi}$, which does not depend on $w$
and its probability space. It can be easily verified that (IRM-
TV-$\ell_{1}$-training-$\Phi$) has the same properties and capabilities as
(IRM-TV-$\ell_{1}$-global-$\Phi$) in Theorem 3.8.

As for (Minimax-TV-$\ell_{1}$), we denote the probability space for $w$
induced by $\rho$ as
$(\Omega_{tr},\mathscr{F}_{\Omega_{tr}},\nu_{tr}^{\rho})$. Following the above
procedure, we obtain
$(\Omega_{all},\mathscr{F}_{\Omega_{tr|all}},\nu_{tr}^{\rho})$ as a coarse
probability space for all $\rho$. Then the following objective function and
optimization model are also well-defined:

$\displaystyle\int_{\Omega_{all}}R(w{\circ}\Phi)\,\mathrm{d}\nu_{tr}^{\frac{1_{(E)}}{E}}{+}\lambda_{\Phi}\max_{\rho}(\int_{\Omega_{all}}|\nabla_{w}R(w{\circ}\Phi)|\,\mathrm{d}\nu_{tr}^{\rho})^{2}=:\mathbb{E}_{w\leftarrow\frac{1_{(E)}}{E}}^{tr}[R(w{\circ}\Phi)]{+}\lambda_{\Phi}\max_{\rho}(\mathbb{E}_{w\leftarrow\rho}^{tr}[|\nabla_{w}R(w{\circ}\Phi)|])^{2},$
(84) $\displaystyle\min_{\Phi}\
\left\\{\mathbb{E}_{w\leftarrow\frac{1_{(E)}}{E}}^{tr}[R(w{\circ}\Phi)]{+}\lambda_{\Phi}\max_{\rho}(\mathbb{E}_{w\leftarrow\rho}^{tr}[|\nabla_{w}R(w{\circ}\Phi)|])^{2}\right\\}.$
(Minimax-TV-$\ell_{1}$-training-$\Phi$)

Similarly, (Minimax-TV-$\ell_{1}$-training-$\Phi$) has the same properties and
capabilities as (Minimax-TV-$\ell_{1}$-global-$\Phi$) in Theorem 3.9.

Although (IRM-TV-$\ell_{1}$-training-$\Phi$) and (Minimax-
TV-$\ell_{1}$-training-$\Phi$) have the same capabilities as (IRM-
TV-$\ell_{1}$-global-$\Phi$) and (Minimax-TV-$\ell_{1}$-global-$\Phi$) to
achieve OOD generalization, the former two are not necessarily equivalent to
the latter two. Denote the probability space for the global environment set as
$(\mathcal{E}_{all},\mathscr{F}_{all},\mu_{all})$ and its induced probability
space for $w$ as $(\Omega_{all},\mathscr{F}_{\Omega_{all}},\nu_{all})$. Let
$\mathscr{F}_{\Omega_{tr\cap
all}}:=\mathscr{F}_{\Omega_{tr|all}}\cap\mathscr{F}_{\Omega_{all}}$ be the
intersection of the two $\sigma$-algebras, which is also a $\sigma$-algebra.
It can be anticipated that $\nu_{tr}(A)=\nu_{all}(A)$, $\forall
A\in\mathscr{F}_{\Omega_{tr\cap all}}$. That is, the measure for both training
and test environments should be the same. However, calculating the integrals
of $R(w{\circ}\Phi)$ and $|\nabla_{w}R(w\circ\Phi)|$ on
$\mathscr{F}_{\Omega_{tr\cap all}}$ may result in distortion when
$\mathscr{F}_{\Omega_{tr\cap all}}\subsetneqq\mathscr{F}_{\Omega_{tr|all}}$.
Because for some $A\in\mathscr{F}_{\mathbb{R}}$, $\\{w:R(w{\circ}\Phi)\in
A\\}\in\mathscr{F}_{\Omega_{tr|all}}\backslash\mathscr{F}_{\Omega_{tr\cap
all}}$ and thus $\mathbb{I}_{\\{w:R(w{\circ}\Phi)\in A\\}}$ cannot be used in
A.1.3 to construct $\int_{\Omega_{all}}R(w{\circ}\Phi)\,\mathrm{d}\nu_{tr}$.

Another similar case is $\Omega_{tr}=\Omega_{all}$ but
$\mathscr{F}_{\Omega_{tr}}\subsetneqq\mathscr{F}_{\Omega_{all}}$. In this
scenario, calculating the integrals of $R(w{\circ}\Phi)$ and
$|\nabla_{w}R(w\circ\Phi)|$ on $\mathscr{F}_{\Omega_{tr}}$ may result in
distortion of (IRM-TV-$\ell_{1}$-global-$\Phi$) and (Minimax-
TV-$\ell_{1}$-global-$\Phi$). This is because for some
$A\in\mathscr{F}_{\mathbb{R}}$, $\\{w:R(w{\circ}\Phi)\in
A\\}\in\mathscr{F}_{\Omega_{all}}\backslash\mathscr{F}_{\Omega_{tr}}$ and thus
$\mathbb{I}_{\\{w:R(w{\circ}\Phi)\in A\\}}$ cannot be used in A.1.3 to
construct $\int_{\Omega_{all}}R(w{\circ}\Phi)\,\mathrm{d}\nu_{all}$.

Last, if $\mathscr{F}_{\Omega_{tr}}\supseteq\mathscr{F}_{\Omega_{all}}$, then
(IRM-TV-$\ell_{1}$-training-$\Phi$) and (Minimax-
TV-$\ell_{1}$-training-$\Phi$) imply (IRM-TV-$\ell_{1}$-global-$\Phi$) and
(Minimax-TV-$\ell_{1}$-global-$\Phi$), as concluded in Corollary 3.12.

##  Appendix B Implementation Details

###  B.1 Subgradient Approach to Solve IRM-TV-$\ell_{1}$ and Minimax-
TV-$\ell_{1}$

First, we give the definition of the Fréchet subdifferential of
$f:\mathbb{R}^{d_{\mathcal{Y}}}\rightarrow\mathbb{R}$ at $w$ denoted by
$\partial f(w)$:

######  Definition B.1 (The Fréchet Subdifferential).

$\displaystyle\partial
f(w){:=}\left\\{v{\in}\mathbb{R}^{d_{\mathcal{Y}}}:\liminf_{\begin{subarray}{c}u\to
w\\\ u\neq
w\end{subarray}}\frac{f(u){-}f(w){-}v\cdot(u{-}w)}{\|u-w\|_{2}}{\geqslant}0\right\\}.$
(85)

The subgradient of the modulus function (i.e., $\ell_{2}$ norm) can be easily
calculated:

$\displaystyle\partial|w|{=}\left\\{\begin{array}[]{ll}\frac{w}{|w|}&\text{if}\quad
w\neq 0,\\\ \\{v{\in}\mathbb{R}^{d_{\mathcal{Y}}}:|v|\leqslant
1\\}&\text{if}\quad w=0.\end{array}\right.$ (88)

In this paper, we directly take $0$ as the convenient subgradient for
$\partial|w|$ at $w=0$, then $\partial|w|$ can be used in a similar way to a
gradient. In particular, the chain rule yields:

$\displaystyle\partial_{\Phi}|\nabla_{w}R(w\circ\Phi)|{=}\left\\{\begin{array}[]{ll}\frac{J_{\Phi}^{\top}[\nabla_{w}R(w\circ\Phi)]*\nabla_{w}R(w\circ\Phi)}{|\nabla_{w}R(w\circ\Phi)|}&\text{if}\quad\nabla_{w}R(w\circ\Phi)\neq
0,\\\ 0&\text{if}\quad\nabla_{w}R(w\circ\Phi)=0,\end{array}\right.$ (91)

where $J_{\Phi}^{\top}[\cdot]$ denotes the transpose of the Jacobian matrix
w.r.t. $\Phi$, and $*$ denotes the matrix multiplication. For a simpler form
with a one-dimensional $w$,

$\displaystyle\partial_{\Phi}|R^{\prime}_{w}(w\circ\Phi)|{=}\mathrm{sign}(R^{\prime}_{w}(w\circ\Phi))\nabla_{\Phi}[R^{\prime}_{w}(w\circ\Phi)].$
(92)

The objective function in (IRM-TV-$\ell_{1}$) and its subgradient are

$\displaystyle g(\Phi)$
$\displaystyle=\mathbb{E}_{w}[R(w\circ\Phi)]+\lambda(\mathbb{E}_{w}[|\nabla_{w}R(w\circ\Phi)|])^{2},\qquad$
(93) $\displaystyle\partial_{\Phi}g(\Phi)$
$\displaystyle=\mathbb{E}_{w}[\nabla_{\Phi}R(w\circ\Phi)]+2\lambda\mathbb{E}_{w}[|\nabla_{w}R(w\circ\Phi)|]\mathbb{E}_{w}[\partial_{\Phi}|\nabla_{w}R(w\circ\Phi)|].$
(94)

Then the update formula for $\Phi$ at the $(k+1)$-th iteration is given by the
following equation, where $\eta>0$ is the learning rate:

$\displaystyle\Phi^{(k+1)}=\Phi^{(k)}-\eta\partial_{\Phi}g(\Phi^{(k)}).$ (95)

We adopt the adversarial training architecture of (Lin et al., 2022) to solve
(Minimax-TV-$\ell_{1}$). Its objective function, gradient and subgradient are

$\displaystyle h(\Phi,\rho)=$
$\displaystyle\mathbb{E}_{w\leftarrow\frac{1_{(E)}}{E}}[R(w{\circ}\Phi)]{+}\lambda(\mathbb{E}_{w\leftarrow\rho}[|\nabla_{w}R(w{\circ}\Phi)|])^{2},$
(96) $\displaystyle\nabla_{\rho}h(\Phi,\rho)=$ $\displaystyle
2\lambda\mathbb{E}_{w{\leftarrow}\rho}[|\nabla_{w}R(w{\circ}\Phi)|]\nabla_{\rho}(\mathbb{E}_{w{\leftarrow}\rho}[|\nabla_{w}R(w{\circ}\Phi)|]),$
(97) $\displaystyle\partial_{\Phi}h(\Phi,\rho)=$
$\displaystyle\mathbb{E}_{w\leftarrow\frac{1_{(E)}}{E}}[\nabla_{\Phi}R(w{\circ}\Phi)]{+}2\lambda\mathbb{E}_{w{\leftarrow}\rho}[|\nabla_{w}R(w\circ\Phi)|]\mathbb{E}_{w{\leftarrow}\rho}[\partial_{\Phi}|\nabla_{w}R(w\circ\Phi)|].$
(98)

Then the update formulae for $\rho$ and $\Phi$ at the $(k+1)$-th iteration are
given by the following equations, where $\eta_{1},\eta_{2}>0$ are the learning
rates:

$\displaystyle\rho^{(k+1)}$
$\displaystyle=\rho^{(k)}+\eta_{1}\nabla_{\rho}h(\Phi^{(k)},\rho^{(k)}),$ (99)
$\displaystyle\Phi^{(k+1)}$
$\displaystyle=\Phi^{(k)}-\eta_{2}\partial_{\Phi}h(\Phi^{(k)},\rho^{(k+1)}).$
(100)

We use the Adam scheme (Kingma & Ba, 2015) for Pytorch333https://pytorch.org
as the optimizer. We apply min-batch subgradients with batch size 1024 in
Landcover, and full-batch subgradients in the other data sets. More
implementing details can be found in the code link, such as the learning rate,
the number of training epochs, etc.

To verify that IRM-TV-$\ell_{1}$ (or Minimax-TV-$\ell_{1}$) will not induce
higher computational complexity than IRM-TV-$\ell_{2}$ (or 2.1/Minimax-
TV-$\ell_{2}$), we run IRM-TV-$\ell_{1}$, Minimax-TV-$\ell_{1}$, IRM-
TV-$\ell_{2}$ and 2.1 for $5300$ epochs and $10$ repetitions, then report the
average running times (in seconds) in the following Table B1. The STDs are all
less than $1.5$ seconds. It indicates that IRM-TV-$\ell_{1}$ (or Minimax-
TV-$\ell_{1}$) has almost the same average running times as those of IRM-
TV-$\ell_{2}$ (or 2.1).

Table B1: Average running times (in seconds) of IRM-TV models by 10 repetitions.  Method | Simulation | House Prices | CelebA | Landcover | Adult
---|---|---|---|---|---
ZIN | 13.5424 | 18.7798 | 23.5480 | 68.0905 | 32.6906
Minimax-TV-$\ell_{1}$ | 12.8399 | 18.7685 | 23.2884 | 68.3897 | 33.1776
IRM-TV-$\ell_{2}$ | 8.9668 | 8.9001 | 22.9603 | N/A | 11.1240
IRM-TV-$\ell_{1}$ | 8.4198 | 9.4552 | 24.0967 | N/A | 11.7834

###  B.2 Synthetic Data Set Generation

We follow (Lin et al., 2022; Tan et al., 2023) to generate synthetic data
sets. Let $t$ be the time in $[0,1]$, and $X_{v}(t)\in\mathbb{R}$ and
$X_{s}(t)\in\mathbb{R}$ be the invariant and spurious features, respectively.
We have

$X_{v}(t)\sim\begin{cases}\mathcal{N}(1,1),~{}\text{w.p.}~{}0.5;\\\
\mathcal{N}(-1,1),\text{w.p.}~{}0.5.\end{cases}Y(t)\sim\begin{cases}\text{sign}(X_{v}(t)),&\text{w.p.}~{}p_{v};\\\
-\text{sign}(X_{v}(t)),&\text{w.p.}~{}1-p_{v}.\end{cases}X_{s}(t)\sim\begin{cases}\mathcal{N}(Y(t),1),&\text{w.p.}~{}p_{s}(t);\\\
\mathcal{N}(-Y(t),1),&\text{w.p.}~{}1-p_{s}(t).\end{cases}$

Moreover, $X_{v}(t)$ and $X_{s}(t)$ are extended to $5$ and $10$ dimensional
sequences by adding up a standard Gaussian noise, respectively.

###  B.3 Architectures of $\Phi$ and $\rho$

We follow (Lin et al., 2022; Tan et al., 2023) to instantiate the Pytorch-
style architectures of the invariant feature extractor $\Phi$ and the
environment inferring measure $\rho$ on each data set, shown in Table B2.
$\Phi$ takes primary variables to learn invariant features, while $\rho$ takes
auxiliary variables to learn environment partition. All the compared methods
except HRM and TIVA share the same $\Phi$, while HRM and TIVA adopt a linear
model for $\Phi$. Minimax-TV-$\ell_{1}$ and 2.1 share the same $\rho$, while
TIVA directly uses auxiliary variables as primary variables in $\Phi$. The
cross entropy and the mean squared error are used as prediction loss functions
for classification tasks and regression tasks, respectively.

Table B2: Pytorch-style architectures of the invariant feature extractor
$\Phi$ and the environment inferring measure $\rho$.

Data Set | $\Phi$ | $\rho$
---|---|---
Simulation | Linear(15, 1) |  Linear(1, 16)$\rightarrow$ReLu()$\rightarrow$Linear(16, 1) $\rightarrow$Sigmoid()
House Prices |  Linear(15, 32)$\rightarrow$ReLu()$\rightarrow$Linear(32, 1) |  Linear(1, 64)$\rightarrow$ReLu()$\rightarrow$Linear(64, 4)$\rightarrow$Softmax()
CelebA |  Linear(512, 16)$\rightarrow$ReLu()$\rightarrow$Linear(16, 1) |  Linear(7, 16)$\rightarrow$ReLu()$\rightarrow$Linear(16, 1)$\rightarrow$Sigmoid()
Landcover |  Conv1d(8, 32, 5)$\rightarrow$ReLu()$\rightarrow$Conv1d(32, 32, 3) | Linear(2, 16)$\rightarrow$ReLu()$\rightarrow$Linear(16, 2)$\rightarrow$Softmax()
$\rightarrow$ReLu()$\rightarrow$MaxPool1d(2, 2)$\rightarrow$Conv1d(32, 64, 3)
$\rightarrow$ReLu()$\rightarrow$MaxPool1d(2, 2)$\rightarrow$Conv1d(64, 6, 3)
$\rightarrow$ReLu()$\rightarrow$AvePool1d(1)
Adult |  Linear(59, 16)$\rightarrow$ReLu()$\rightarrow$Linear(16, 1) |  Linear(6, 16)$\rightarrow$ReLu()$\rightarrow$Linear(16, 4)$\rightarrow$Softmax()

##  Appendix C Invariant Feature Identifiability

We take the simulation study with
$(p_{s}^{-},p_{s}^{+},p_{v}(t))=(0.999,0.9,0.8)$ as an example. There are $15$
features in this experiment, where No. 1$\sim$5 are invariant features and No.
6$\sim$15 are spurious features. We run $10$ times of this experiment with
different simulated samples and compute the normalized average absolute values
of the feature weights, shown in Table C1 and Figure C1. The invariant
features of IRM-TV-$\ell_{1}$ and Minimax-TV-$\ell_{1}$ take up $64.88\%$ and
$65.08\%$ of the total feature weights, respectively. In contrast, the
invariant features of IRM-TV-$\ell_{2}$ and 2.1 take up only $58.88\%$ and
$55.90\%$ of the total feature weights, respectively. Hence the TV-$\ell_{1}$
models extract more invariant features than the TV-$\ell_{2}$ models.
Moreover, the gap between the invariant and spurious feature weights in
Figures C1(b) or C1(d) is larger than that in Figures C1(a) or C1(c),
respectively. For example, the invariant feature weight w2$=0.0884$ and the
spurious feature weights w6$=0.0679$ and w8$=0.0586$ for 2.1, while
w2$=0.1210$, w6$=0.0403$, and w8$=0.0446$ for Minimax-TV-$\ell_{1}$,
respectively. It indicates that Minimax-TV-$\ell_{1}$ enhances the invariant
feature w2 while suppresses the spurious features w6 and w8.

Table C1: Normalized absolute values of feature weights for IRM-TV models in
simulation study with $(p_{s}^{-},p_{s}^{+},p_{v}(t))=(0.999,0.9,0.8)$.
w1$\sim$w5 correspond to invariant features and w6$\sim$w15 correspond to
spurious features.

Method | w1 | w2 | w3 | w4 | w5 | w6 | w7 | w8 | w9 | w10 | w11 | w12 | w13 | w14 | w15
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---
IRM-TV-$\ell_{2}$ | 0.1091 | 0.1143 | 0.1452 | 0.0912 | 0.1290 | 0.0412 | 0.0439 | 0.0454 | 0.0323 | 0.0483 | 0.0361 | 0.0438 | 0.0282 | 0.0526 | 0.0395
IRM-TV-$\ell_{1}$ | 0.1232 | 0.1292 | 0.1548 | 0.1090 | 0.1326 | 0.0329 | 0.0329 | 0.0446 | 0.0249 | 0.0371 | 0.0317 | 0.0386 | 0.0262 | 0.0482 | 0.0341
ZIN | 0.1253 | 0.0884 | 0.1118 | 0.1336 | 0.0999 | 0.0679 | 0.0454 | 0.0586 | 0.0353 | 0.0261 | 0.0267 | 0.0376 | 0.0510 | 0.0505 | 0.0419
Minimax-TV-$\ell_{1}$ | 0.1320 | 0.1210 | 0.1410 | 0.1353 | 0.1215 | 0.0403 | 0.0356 | 0.0446 | 0.0221 | 0.0279 | 0.0285 | 0.0341 | 0.0352 | 0.0395 | 0.0412

(a) IRM-TV-$\ell_{2}$ (b) IRM-TV-$\ell_{1}$

(c) ZIN (d) Minimax-TV-$\ell_{1}$

Figure C1: Normalized absolute values of feature weights for IRM-TV models in
simulation study with $(p_{s}^{-},p_{s}^{+},p_{v}(t))=(0.999,0.9,0.8)$.
w1$\sim$w5 correspond to invariant features and w6$\sim$w15 correspond to
spurious features.