# Learning Optimal Features via Partial Invariance

Moulik Choraria1,111corresponding author, Ibtihal Ferwana1, Ankur Mani2, Lav
R. Varshney1

###### Abstract

Learning models that are robust to distribution shifts is a key concern in the
context of their real-life applicability. Invariant Risk Minimization (IRM) is
a popular framework that aims to learn robust models from multiple
environments. The success of IRM requires an important assumption: the
underlying causal mechanisms/features remain invariant across environments.
When not satisfied, we show that IRM can over-constrain the predictor and to
remedy this, we propose a relaxation via _partial invariance_. In this work,
we theoretically highlight the sub-optimality of IRM and then demonstrate how
learning from a partition of training domains can help improve invariant
models. Several experiments, conducted both in linear settings as well as with
deep neural networks on tasks over both language and image data, allow us to
verify our conclusions.

##  1 Introduction

Standard machine learning models trained using classical Empirical Risk
Minimization (ERM) can be expected to generalize well to unseen data drawn
from the same distribution as the training data (Vapnik 2013). However,
distribution shifts during test time (when data is from different sources or
under different conditions) can severely degrade model performance (Beery, V.
Horn, and Perona 2018; Lake et al. 2017; Marcus 2018). The errors can often be
attributed to the model picking up statistically informative but spurious
correlations, which in turn limits their real-life applications since in
practice, the use-case almost always differs from the training data. Thus,
several lines of research explore alternate learning objectives for training
robust models.
One particular line of research stems from the Invariant Causal Prediction
framework (Peters, Bühlmann, and Meinshausen 2015), where the goal is to learn
causal mechanisms that work well under interventions; our work focuses on the
similarly inspired Invariant Risk Minimization (IRM) framework, which aims to
learn a predictor that relies only on features that are invariant across all
training environments. The underlying motivation for invariance is rooted in
its strong links with causality (Pearl 2009), with the intuition being that by
invariance can help the model distinguish the causal features from domain-
specific spurious features, which it can then discard for better
generalization.
A standard assumption in such invariance-based objectives is that of
sufficiency (Ahuja et al. 2020b): there exists a predictor, relying solely on
the invariant features, which can achieve Bayes optimal risk in all
environments. While fairly general, this assumption may not be satisfied for
certain classes of distribution shifts. For instance, consider a prediction
task with concept drift, wherein the relationship of the ‘causal’ features
(features that are causally responsible to the label) with the label changes
across training environments. Here, a predictor relying solely on invariant
features ends up being over-constrained, since it is incentivized to discard
non-invariant but informative causal features that are needed for Bayes
optimality. Such situations are ubiquitous in practice, for instance in
language tasks in which linguistic features can have different connotations
within different communities (Gallacher 2021; Mani, Varshney, and Pentland
2021) or in tasks with distribution shifts across time (Luu et al. 2021).
Additionally, even when a sufficient representation exists theoretically, it
may not be accessible due to shortcomings in the optimization of the learning
objective. However, these factors are seldom accounted for when considering
the application of IRM (or other invariant learning objectives) for a given
use-case (Peyrard et al. 2021; Adragna et al. 2020). Thus, it is important to
develop a characterization for the same.
To address this gap in literature, we present a first study to characterize
the behaviour of IRM under explicit concept drifts. Then, we take a step
further and propose a relaxation for IRM via the _Partial Invariance_ (P-IRM)
framework. We find that our framework increases the flexibility of invariant
models by allowing learning of features that are locally invariant within a
partition of the training environments. This flexibility is accompanied with
an inherent trade-off; the cost of finding the right partition, in an
information-agnostic setting, grows exponentially in the number of
environments. However, for certain classes of problems, including the language
tasks alluded to previously, readily available meta-information often allows
us to easily infer the ‘optimal’ training partition for a given use-case.
Notice that in doing so, we move away from the OoD minimax regime, and instead
focus on improving generalization conditioned on availabity of this meta-
information. In this work, we begin by presenting a theoretical
characterization of IRM under concept shifts. Next, we formally quantify the
notion of meta-information and assuming access to it, we theoretically and
empirically demonstrate how the notion of Partial Invariance can help improve
the performance of invariant models. The rest paper is organized as follows:
we begin with a literature review in Sec. 2, and motivate P-IRM and present
our main results in Sec. 3. We report our empirical evaluations in Sec. 4 and
wrap up with some concluding remarks in Sec. 5.

##  2 Related Work

Many approaches aim to learn deep invariant feature representations: some
focus on domain adaptation by finding a representation whose distribution is
invariant across source and target distributions (Ben-David et al. 2010;
Zhang, Gong, and Schoelkopf 2015), while others focus on conditional domain-
invariance (Gong et al. 2016; Li et al. 2018). However, there is evidence that
domain adaption approaches are insufficient when the test distribution may lie
outside the convex hull of training distributions (Lee and Raginsky 2018;
Duchi, Glynn, and Namkoong 2021; Mohri, Sivek, and Suresh 2019). Other
approaches include Bayesian Deep Learning (Neal 1996), which tries to account
for model uncertainty during test-time, and Robust Optimization (Ben-Tal, El
Ghaoui, and Nemirovski 2009), which aims to generalize well to distributions
close to training.
Our work focuses particularly on the IRM framework (Arjovsky et al. 2019),
which relates to domain generalization wherein access to the test distribution
is not assumed. IRM is rooted in the theory of causality (Schölkopf et al.
2012) and proposes invariance for achieving OoD generalization (Peters,
Bühlmann, and Meinshausen 2016; Heinze-Deml, Peters, and Meinshausen 2018). In
(Ahuja et al. 2020a), the authors reformulate IRM via a game-theoretic
approach, wherein the invariant representation corresponds to the Nash
equilibrium of a game. While the IRM framework assumes only the invariance of
the conditional expectation of the label given the representation, some
follow-ups rely on stronger invariance assumptions (Xie et al. 2021; Mahajan,
Tople, and Sharma 2021). As mentioned before, this line of work assumes
sufficiency of invariant features whereas we specifically focus on
distribution shifts when sufficiency is violated.
Several follow-up works attempt to characterize IRM’s performance under
different settings and model assumptions. It has been noted that carefully
tuned ERM can often outperform state-of-the-art domain generalization
approaches, including IRM, across multiple benchmarks (Gulrajani and Lopez-Paz
2020). The failure of IRM may stem from the gap between the proposed framework
and its practical “linear” version (IRMv1), which fails to capture natural
invariances (Kamath Pritish and Srebro 2021). Indeed, the authors of
(Rosenfeld, Ravikumar, and Risteski 2020) demonstrate that a near-optimal
solution to the IRMv1 objective, which matches IRM on training environments,
does no better than ERM on environments that differ significantly from
training. Following these deficiencies, several works propose alternate
objectives for achieving invariance (Krueger et al. 2021; Bellot and van der
Schaar 2020; Jin, Barzilay, and Jaakkola 2020; Ahuja et al. 2021; Shui, Wang,
and Gagné 2021).
However, unlike previous works that aim to improve the invariance learning
objective, we question whether invariance as a constraint can be improved upon
for better performance. To that end, our notion of _partial invariance_
generalizes not only IRM, but all similar invariance learning objectives. The
use of meta-information for invariant learning has been proposed in (Lin, Zhu,
and Cui 2022). However, unlike partitioning, the focus therein is to
artificially generate environment membership for samples when not available a
priori. Finally, a related idea appears in (Yu et al. 2022), which proposes
applying different invariance penalty weights for different domains, but with
the goal of addressing data quality variance across domains.

##  3 Theory

In this section, we present the notion of partial invariance.
Notation: We use upper-case boldface $\bm{U}$ to denote matrix/tensor/vector
valued random variables, and lowercase boldface $\bm{u}$ to denote scalar
valued random variables. We use upper-case $U$ to denote
matrices/vectors/tensors and lowercase $u$ to denote scalars.

###  3.1 Invariant Risk Minimization

The IRM setup assumes access to datasets of the form
$D_{e}:=\\{X^{e}_{i},y^{e}_{i}\\}_{i=1}^{n_{e}}$ collected from multiple
training environments $e\in\mathcal{E}_{tr}$. The samples in dataset $D_{e}$
are i.i.d. from the environment’s joint distribution,
$P(\bm{X}^{e},\bm{y}^{e})$. The task is to estimate a map
$f:\mathcal{X}\rightarrow\mathcal{Y}$ or alternatively, the conditional
distribution $P(Y|X)$, so that it performs well across unseen environments
$\mathcal{E}_{all}\supset\mathcal{E}_{tr}$. Formally, the IRM framework aims
to minimize the Out-of-Distribution (OoD) risk:
$R^{OoD}(f)=\max_{e\in\mathcal{E}_{all}}R^{e}(f)$, where
$R^{e}(f):=\bm{E}_{\bm{X}^{e},\bm{y}^{e}}[\ell(f(\bm{X}^{e}),\bm{y}^{e})]$ is
the expected risk in environment $e$. The predictor $f$ is parametrized as
$w\circ\Phi$, wherein $\Phi:\mathcal{X}\rightarrow\mathcal{Z}$ represents the
learned representation and $w:\mathcal{Z}\rightarrow\mathcal{Y}$ is a linear
predictor over said representation. The IRM learning objective is posed as a
constrained optimization problem:

$\displaystyle\min_{\Phi,w}\sum_{e\in\mathcal{E}_{obs}}{R}^{e}(w\circ\Phi)$
(IRM) s.t. $\displaystyle\
w\in\text{arg}\min_{\tilde{w}}{R}^{e^{\prime}}(\tilde{w}\circ\Phi)\ \forall\
{e^{\prime}}\in\mathcal{E}_{tr}\mbox{.}$ (1)

To avoid the inner optimization, the minimization constraint is replaced by a
more tractable gradient penalty:

$\displaystyle\min_{\Phi,w}\sum_{e\in\mathcal{E}_{obs}}{R}^{e}(w\circ\Phi)$
(IRMgc) s.t. $\displaystyle\
w\in\\{{\tilde{w}}:\|\nabla_{w}{R}^{e^{\prime}}(w\circ\Phi)\|=0\ \forall\
{e^{\prime}}\in\mathcal{E}_{tr}\\}\mbox{,}$ (2)

where IRMgc is shorthand for the gradient constrained IRM. In practice, this
constraint is enforced via a regularizer $\lambda$:

$\displaystyle\min_{\Phi}\sum_{e\in\mathcal{E}_{obs}}{R}^{e}(\Phi)+\lambda\|\nabla_{w,w=1.0}{R}^{e}(\Phi)\|\mbox{,}$
(IRMv1)

where the implicit overparametrization in having a separate classifier and
representation map is removed by fixing a dummy classifier $w=1.0$. Thus,
$\Phi$ becomes the entire invariant predictor and the strictness of the
gradient norm penalty, which enforces invariance, is via $\lambda$. Note that
when $\lambda=\infty$, IRMv1 is equivalent to IRMgc, which in turn is the
first order approximation for the true IRM objective.
An intrinsic assumption in the IRM learning setup for proving minimax
optimality is the ideal scenario of sufficiency i.e. there exists a $\Phi$
that is invariant across all $e\in\mathcal{E}_{tr}$ and is sufficient i.e
$\bm{y}^{e}\perp\bm{X}^{e}|\ \Phi(\bm{X}^{e})\ \forall e\in\mathcal{E}_{all}$
(Ahuja et al. 2020b). However if sufficiency is violated for an environment,
one would expect the IRM model, which relies solely on invariant features, to
be sub-optimal for that environment (compared to a model that utilizes non-
invariant features along with invariant ones). Such a situation may arise
under _concept drift_ , wherein the the conditional expectation of the label
$\bm{y}_{e}$ given “causal” features may change across environments. Thus, in
practice, if an invariant $\Phi$ that is also sufficient for environments does
not exist for the desired use-case, we expect the performance of IRM (or
related frameworks) to degrade. We illustrate this with a simple example.
Example 1 We adapt the generative model from (Arjovsky et al. 2019): the goal
is to predict target $y$ using $X=[x_{1},x_{2},x_{3}]$, in environment $e$
such that $e\in\mathcal{E}_{all}$ can affect the distribution of $X$ as well
as the conditional distribution of $y$ given $X$ via a deterministic map
$c(e):\mathcal{E}_{all}\rightarrow\\{-1,1\\}$:

$\begin{gathered}\bm{x}_{1}\leftarrow N(0,\sigma(e)^{2}),\bm{x}_{2}\leftarrow
N(0,\sigma(e)^{2}),\\\ {c(e)\in\\{1,-1\\},\epsilon\sim N(0,\sigma(e)^{2})}\\\
\bm{y}\leftarrow\bm{x}_{1}+c(e)\bm{x}_{2}+\epsilon,\ \epsilon\perp\bm{x}_{1},\
\epsilon\perp\bm{x}_{2}\\\
\bm{x}_{3}\leftarrow\bm{y}+N(0,1),\sigma(e)^{2}\in[0,\sigma_{MAX}^{2}].\\\
\end{gathered}$

We estimate $\bm{y}$ as
$\hat{\bm{y}}=\alpha_{1}\bm{x}_{1}+\alpha_{2}\bm{x}_{2}+\alpha_{3}\bm{x}_{3}$.
Within the IRM framework, the only feasible representation $\Phi$ (upto
scaling) that yields invariant predictors across all $e$ is
$\Phi([\bm{x}_{1},\bm{x}_{2},\bm{x}_{3}])=[\bm{x}_{1},0,0]$, with
corresponding regression coefficients $[1,0,0]$. Although this minimizes the
OoD error for arbitrary $e$, it does so by discarding the non-invariant but
informative $\bm{x}_{2}$. However, if our predictor is privy to some knowledge
of $c(e)$, we could first partition the set of training environments
$\mathcal{E}_{tr}$ into two partitions, such that environments within a
partition have the same $c(e)$ value. Then, applying IRM within each partition
yields models with better performance that can exploit $\bm{x}_{2}$ as an
invariant feature in the partition. Note that this partial notion of
invariance still retains the ability to discard spurious/non-causal
$\bm{x}_{3}$. Additionally with partitioning, we can improve generalization if
information about $c(e_{unseen})$ is available, by choosing the right
model/partition for prediction. Next, we study the conditions under which
partitioning can improve upon IRM performance and we refer to this method as
P-IRM.

###  3.2 Model

For our analysis, we consider a simple regression task to succinctly capture
our intuition about the conditions under which partitioning is feasible. To
begin with, we assume access to the underlying causal features and instead,
focus on understanding the nature of the IRM solution set under distribution
shifts. In the next part, we extend this analysis to study learning under
partial invariance.
We consider the following generative model: we observe samples
$(X_{i}^{e},y_{i}^{e})$in environment $e$ and the goal is to predict
$y_{i}^{e}$ from $X_{i}^{e}$. $X_{i}^{e}$’s are samples corresponding to the
random variable $\bm{X}^{e}\sim P(\bm{X}^{e})$, as described below:

$\bm{X}^{e}=[\bm{x}_{1}^{e},\bm{x}_{2}^{e},\ldots,\bm{x}_{c}^{e}]^{\top}\sim
P(\bm{X}^{e}),$

where each $\bm{x}_{i}^{e}$ denotes an individual feature. To simplify our
initial analysis, we assume that the individual features are independent of
each other and are normalized i.e. $\bm{E}[\bm{X}^{e}]=\bm{0}$ and
$\bm{E}[\bm{X}^{e}{\bm{X}^{e}}^{\top}]=\bm{I}\ \forall\ e$. The target
$\bm{y}^{e}$ for given $\bm{X}^{e}$ can be characterized as:

$\displaystyle\begin{split}&\bm{y}^{e}=\langle{W^{e}},\bm{X}^{e}\rangle+\epsilon_{y}\mbox{,}\\\
&W^{e}=[w_{1}^{e},w_{2}^{e},\ldots,w_{c}^{e}]\in\mathbb{R}^{c},\epsilon_{y}\sim
N(0,\sigma_{y}^{2}(e)).\end{split}$ (3)

where weights $W^{e}$ encode the conditional distribution of observing label
$\bm{y}^{e}$ given $\bm{X}^{e}$ in environment $e$ and are fixed for that
environment, and $\langle\cdot,\cdot\rangle$ denotes the standard inner
product in $\mathbb{R}^{c}$. For a given feature $\bm{x}_{i}^{e}$ in
environment $e$, the corresponding feature weight $w_{i}^{e}$ is independently
and uniformly sampled from set $A_{i}$ for each environment $e$. Once sampled
however, these weights remain fixed for that environment. Additionally
$|A_{1}|=1$, so that feature weight $w_{1}^{e}$ is fixed and thus
$\bm{x}_{1}^{e}$ is invariant for all $e$:

$\begin{gathered}\bm{W}^{e}=[\bm{w}_{1}^{e},\bm{w}_{2}^{e},\ldots,\bm{w}_{c}^{e}],\\\
\text{ where }\ {\bm{w}_{i}}^{e}\sim\text{Unif}(\\{A_{i}\\})\ \forall\
i\in\\{1,2,\ldots,c\\},\\\ |A_{1}|=1,\ |A_{i}|>1\ \forall\ i>1.\end{gathered}$
(4)

We make note of some important aspects. As per our model, ${\bm{x}_{1}^{e}}$
is the only _truly invariant feature_ since
$\bm{E}[\bm{y}^{e}|{\bm{x}_{1}^{e}}]=w_{inv}.\bm{x}_{1}^{e}$, is fixed for all
$e$, where $A_{1}=\\{w_{inv}\\}$ is a singleton, and $w_{inv}$ denotes the
invariant feature weight. Additionally, the cardinality of set $A_{i},|A_{i}|$
defines an implicit notion of the variance of feature $x_{i}$, with a higher
cardinality indicating that the feature weight is more likely to change across
environments and is thus, _less invariant_.
With our generative model in place, we next consider the task of predicting
$\bm{y}^{e}$ given $\bm{X}^{e}$, under the mean squared loss. Recall that the
IRM framework considers predictors of the form $w\circ\Phi$, where the
transformation $\Phi$ extracts a suitable representation and $w$ is the linear
predictor acting on that representation. Due to the implicit
overparametrization, we fix $w=1.0$ to a scalar value as proposed in (Arjovsky
et al. 2019) and analyze the corresponding IRM solutions with
${\Phi}\in\mathbb{R}^{c}$. For simplicity, we ignore finite sample effects and
consider the objective in (IRMv1) when $\lambda=\infty$, or equivalently, the
gradient penalty constraint in (2) which ideally approximates the true IRM
objective. Additionally, we assume the following for training environments
$\mathcal{E}_{tr}$.

######  Assumption 1 (_Sufficiency for IRM_).

Assume $\exists$ an environment $e\in\mathcal{E}_{tr}$ for which the truly
invariant predictor is sufficient, i.e. the corresponding feature weights
satisfy $w_{1}^{e}=w_{inv}$ and ${w^{e}}_{i}=0\ \forall\
i\in\\{2,\ldots,c\\}$.

In other words, we assume existence of a training environment in which the
invariant predictor that only recovers the invariant feature $\bm{x}_{1}^{e}$
achieves optimal MSE risk, which is a standard assumption in related
literature (Ahuja et al. 2020b).

######  Lemma 1.

As per above parametrization (with $w=1.0$), under Assumption 1, the values
for ${\Phi}$ that satisfies the IRM solution constraints in (2) is a singleton
and the value of the corresponding predictor equates to
${\Phi}=[w_{inv},0,0\ldots,0],$ the predictor only recovers the invariant
feature.

The proof of the lemma is included in the Appendix and relies on showing that
any predictor which assigns non-zero weights to any of non-invariant features
would violate the gradient penalty constraints. More importantly, the previous
lemma roughly says that any non-invariant feature will be discarded by the IRM
predictor. Note that while this is a desirable property for minimax
optimality, we ask whether we can do better given additional contextual
information.
We formalize the notion of contextual information explicitly by defining an
oracle $\omega(e)=\mathbf{1}[\|W^{e_{ref}}-W^{e}\|_{0}\leq\delta]$, that
provides us a notion of distance between environments, from a fixed reference
environment $e_{ref}$. Alternatively, it identifies whether environment $e$ is
close to $e^{ref}$.
Remark 1: The choice of the $\ell_{0}$ metric for the oracle suits our
combinatorial setting, since we do make any assumptions on the individual
elements in the feature weight sets (i.e. $A_{i}$’s).
Next we characterize our objective to utilize this information. Suppose we
know that our test environment shares the feature weight with the reference
environment for a given feature $\bm{x}_{i}^{e}$. Then we can define the goal
of minimizing the risk w.r.t. to the predictor $f$, conditioned on this
information:

$R^{cond}(f)=\bm{E}_{e\mbox{ s.t. }w^{e_{ref}}_{i}=w^{e}_{i}}R^{e}(f),$

where the expectation is over the draw of environments as per the uniform
sampling. We note that a predictor that accounts for the prior condition
(reference feature) will improve performance (i.e. with a lower MSE risk
$R^{cond}$), as compared to the truly invariant predictor in the previous
lemma. However, to obtain the required feature as a feasible solution via IRM
constraints, we need to first isolate a subset of training environments
$\mathcal{E}_{partition}\subseteq\mathcal{E}_{tr}$ such that within this set,
$w^{e}_{i}$ is invariant and secondly, that we avoid learning the rest of the
non-invariant features to avoid feature weight mismatches in unseen
environments. It turns out that with access to the oracle and under certain
mild conditions, we can ensure exactly that in our uniform distribution shift
model. Before stating the result, we require a similar sufficiency assumption
for the _partially invariant predictor_.

######  Assumption 2 (Sufficiency for P-IRM).

Assume $\exists$ an environment $e\in\mathcal{E}_{tr}$ for which the partially
invariant predictor is sufficient, i.e. the corresponding feature weights
satisfy $w_{1}^{e}=w_{inv}$, $w_{i}^{e}=w_{i}^{e_{ref}}$ and ${w^{e}}_{j}=0\
\forall\ j\in\\{2,\ldots,c\\}\setminus\\{i\\}$.

######  Theorem 1.

Under the model (4), under Assumption 2, with access to oracle
$\omega(e)=\mathbf{1}[\|W^{e_{ref}}-W^{e}\|_{0}\leq\delta]$ and
$\delta<(c-2)/2$, isolate
$\mathcal{E}_{partition}:=\\{e\in\mathcal{E}_{tr}|\omega(e)=1\\}\cup\\{e_{ref}\\}\subseteq\mathcal{E}_{tr}$.
Next, let $|A_{i}|=k$, where $A_{i}$ is the set corresponding to the feature
weight $w_{i}^{e_{ref}}$ of interest. Then, if the sets $\\{A_{j}\\}\forall\
j\in\\{2,\ldots,c\\}\setminus\ \\{i\\}$ satisfy $|A_{j}|>\alpha k$ for some
$\alpha>1$, we have with probability greater or equal to
$(\frac{p}{p+1})^{|\mathcal{E}_{partition}|}$, where
$p\geq\frac{(c-1-\delta)\alpha}{\delta}$, the IRM solution over set
$\mathcal{E}_{partition}$ will recover the feature of interest
$w_{i}^{e_{ref}}$.

The proof, available in the Appendix, relies on showing that within the
partition that satisfies the oracle condition, the probability of successfully
isolating the required feature is high. Then the result follows as a
consequence of Lemma 1.
In words, the theorem says that if we can identify a partition in which the
environments are not too different, then with high probability, the IRM
solution will recover features which do not vary too much (i.e. non-invariant
but still close to invariant). Note that in case of erroneous partitioning,
the solution set allowed by the non-convex penalty becomes harder to
characterize due to the presence of other feature weights besides the
reference. Nevertheless, if the conditions are such that probability of that
happening is sufficiently low, we can safely assume that partitioning will
achieve a better expected risk. Additionally, it suggests that P-IRM becomes
feasible as the oracle becomes more precise and the feature of interest is
close to invariant.
Remark 2: While P-IRM does improve upon the IRM solution, both variants are
likely to be outperformed by ERM in this setting. However, we point out that
this is a simplified setting wherein access to causal features is assumed. In
more general settings when the causal features need to be inferred from
complex data, ERM may be susceptible invariance to confounders/anti-causal
variables and thus, we require invariance as a means to make the solution
robust.

###  3.3 Partitioning and Partial Invariance

Next, we study P-IRM in a general setup, using previous results to
characterize the required number of training environments as in IRM. As
before, we assume access to the oracle, $\omega$ to identify the partition,
i.e. $\mathcal{E}_{partition}\subseteq\mathcal{E}_{tr}$.
Learning Setup: We consider the same causal mechanism for regression task
$(\bm{x}^{e},y^{e})$ from before. The goal is to find a partition using the
oracle such that a feature of interest corresponding to the reference
environment, $w_{i}^{e_{ref}}$ is retained. Note that since we want to retain
only the invariant features denoted as
${\bm{X}_{inv}}^{e}=[\bm{x}_{1}^{e},\bm{x}_{i}^{e}]$, and discard the non-
invariant (or non-partially invariant) features, we encapsulate them into the
noise term as $\tilde{\epsilon}_{y}={\epsilon}_{y}+(\bm{X}^{e}_{\\{1\cdots
c\\}\setminus\\{1,i\\}})^{\top}W^{e}_{\\{1\cdots c\\}\setminus\\{1,i\\}}$.
Then, notice that we still have $\tilde{\epsilon}_{y}\perp{\bm{X}_{inv}}^{e}$
and that $\bm{E}[\tilde{\epsilon}_{y}]=0$, due to feature independence and
centering assumptions. Next, we consider a realistic learning setup where we
observe a scrambled version $\tilde{\bm{X}}^{e}$ of the true causal features
${\bm{X}^{e}}$:

$\begin{gathered}\bm{y}^{e}=({\bm{X}_{inv}}^{e})^{\top}{W_{inv}}^{e}+{\tilde{\epsilon}}_{y},\
\ \tilde{\epsilon}_{y}\perp{\bm{X}_{inv}}^{e},\ \
\bm{E}[\tilde{\epsilon}_{y}]=0\\\
\tilde{\bm{X}}^{e}=S(\bm{X}^{e},\bm{X^{\prime}}^{e}).\end{gathered}$ (5)

Here, $\bm{X}^{e}=[{\bm{X}_{inv}}^{e},\bm{X}^{e}_{\\{1\cdots
c\\}\setminus\\{1,i\\}}]\in\mathbb{R}^{c}$ denote the causal features with
respect to the label, $\bm{X^{\prime}}^{e}\in\mathbb{R}^{q}$,
$\tilde{\bm{X}}^{e}=S(\bm{x}^{e},\bm{d}^{e})\in\mathbb{R}^{d}$ with
$S\in\mathbb{R}^{d\times(c+q)}$. The variable $\bm{X^{\prime}}^{e}$ may be
arbitrarily correlated with ${\bm{X}_{inv}}^{e},\tilde{\epsilon}_{y}$ or the
label $\bm{y}^{e}$ and is intended to represent the spurious correlations in
data. However, we require $S$ to be such that $\exists\tilde{S}$ s.t.
$\tilde{S}(S(\bm{X}^{e},\bm{X^{\prime}}^{e}))={\bm{X}_{inv}}^{e}$ i.e. an
inverse map such that the recovery of the desired features is feasible.
Next, we define $\gamma=\frac{1}{k\sqrt{2n}}\exp(-nD(\delta/n\|1/\alpha k)),$
where as before, $\delta$ is the oracle distance parameter, $k$ is the
cardinality of the set $A_{i}$, $|A_{i}|$, $\alpha$ is as defined in Theorem
1, $n=c-2$ and $D(m\|n)$ denotes KL divergence between $\mbox{Bern}(m)$ and
$\mbox{Bern}(n)$. Intuitively, $\gamma$ estimates the lower bound on the
probability of sampling an environment under the generative model that
satisfies the oracle condition of close distance to the reference environment.
Then we have the following sample complexity on the number of required
environments.

######  Theorem 2 (Informal).

Assume we observe $(\tilde{\bm{X}}^{e},\bm{y}^{e})$ as per (5), with
environments $e\in\mathcal{E}_{tr}$ sampled as per (4) and let
$\mathcal{E}_{partition}:=\\{e\in\mathcal{E}_{tr}|\omega(e)=1\\}\cup
e_{ref}\subseteq$. Let $\bm{\Phi}\in\mathbb{R}^{d\times d}$ have rank $r>0$.
Then sampling $|\mathcal{E}_{tr}|>\frac{1}{\gamma}(d-r+d/r)\log(1/\epsilon)$
ensures partition cardinality $|\mathcal{E}_{partition}|>d-r+d/r$ with
probability $>1-\epsilon$. Furthermore, if $e\in\mathcal{E}_{partition}$ lie
in linear general position of degree $r$ (Assumption 3 in Appendix), then with
probability greater than or equal to
$(\frac{p}{p+1})^{|\mathcal{E}_{partition}|}$, where
$p\geq\frac{(c-1-\delta)\alpha}{\delta}$, the oracle identifies
$\mathcal{E}_{partition}$ such that the predictor $w\circ\Phi$ learnt via IRM
within that partition recovers the desired features/weights and corresponding
prediction $({\bm{X}}_{inv}^{e})^{\top}{W}_{inv}^{e}$, $\forall
e\in\mathcal{E}_{all}$ which satisfy $w^{e}_{i}=w^{e_{ref}}_{i}$.

The proof along with the formal statement is included in the Appendix and
follows from our previous results by applying concentration bounds on the draw
of environments, and subsequently using prior generalization results for IRM.
In words, Theorem 2 states that if the obtained partition is accurate, is of
sufficient cardinality and is sufficiently diverse, then ${\Phi}$ recovers the
partially invariant features. However, notice that the required number of
environments grows inversely with $\gamma$, meaning that we need stronger
priors (i.e. sample environments close to the reference) to obtain feasible
sample complexities in the number of required environments.

###  3.4 Partial Invariance in Practice

Next, we state the P-IRM objective more formally. We first assume a distance
metric $d$ between environments (known directly or via contextual
information). Then, our goal is to identify a subset of training environments
$\mathcal{E}_{partition}\subseteq\mathcal{E}_{tr}$ such that its average
distance w.r.t. a reference environment $e^{ref}$ roughly satisfies:

$\frac{1}{|\mathcal{E}_{partition}|}\sum_{e\in\mathcal{E}_{partition}}d(e,e^{ref})<\frac{1}{|\mathcal{E}_{tr}|}\sum_{e\in\mathcal{E}_{tr}}d(e,e^{ref}).$

Thus, the predictor is trained on a subset of observed environments. However,
discarding environments is not data-efficient and can lead to lower fidelity
and worse generalization, especially in high-complexity models. To avoid this,
we introduce the notion of conditional invariance as an alternative. Formally,
consider the set of observed training environments $\mathcal{E}_{tr}$ and a
subset corresponding to the partition $\mathcal{E}_{partition}$ (chosen
suitably via $d$), satisfying
$\mathcal{E}_{partition}\subseteq\mathcal{E}_{tr}$. We propose the following
two variants of P-IRM:

$\displaystyle\min_{\Phi,w}\sum_{e\in\mathcal{E}_{1}}{R}_{e}(w\circ\Phi)\mbox{
s.t.}\ w\in\text{arg}\min_{\tilde{w}}{R}_{e^{\prime}}(\tilde{w}\circ\Phi)\
\forall\ {e^{\prime}}\in\mathcal{E}_{2}\mbox{,}$ $\displaystyle\ \text{if
}\mathcal{E}_{1}=\mathcal{E}_{2}=\mathcal{E}_{partition},$ (P-IRM
(Partitioning)) $\displaystyle\ \text{if }\mathcal{E}_{1}=\mathcal{E}_{tr}\
\text{\&}\ \mathcal{E}_{2}=\mathcal{E}_{partition}$ (P-IRM (Conditioning))

where the empirical risk minimization objective is over environments in
$\mathcal{E}_{1}$ and the IRM invariance constraint is applied on environments
in $\mathcal{E}_{2}$. For P-IRM (Conditioning), note that while the model uses
data from all environments, the invariance penalty is applied only to
environments within the chosen partition, which mitigates the issue of having
fewer data samples. Intuitively, it serves as a relaxation of the IRM
objective to allow for partially invariant features. Next, we qualitatively
discuss some potential issues in the application of P-IRM. Firstly, fulfilling
the requirements as per Theorem 2, for the required worst case number of
environments is infeasible. Fortunately, in practice, IRM can pick up the
required invariances from just two environments and we expect P-IRM to
overcome that issue as well.
Next, we revisit the idea of the distance oracle. While a precise
characterization of the distance between causal features of different domains
is essentially unobtainable in practice, certain situations allow for
inferring the nature of the distribution shift via available contextual
information which, while often discarded by practictioners, can serve as an
effective pseudo-metric for the same. For instance, authors of (Luu et al.
2021) pointed out that temporal mis-alignments of distributions in language
tasks leads to performance degradation, noting that degradation increases with
an increase in the time duration between test and train environments. Thus,
learning from only the recent past could yield a larger and more relevant set
of invariant features for a use-case on future data.

##  4 Experiments

We start with a basic sanity check via a synthetic experiment, as an extension
of the example presented earlier to visualize how IRM can end up suppressing
non-invariant causal features, leading to performance degradation. We then
evaluate the efficacy of the P-IRM framework (both partitioning or
conditioning) on four tasks: a regression task for housing price prediction,
an image classification task on the MetaShift dataset (Liang and Zou 2022), an
entity recognition task for scientific texts on the SciERC dataset (Luan et
al. 2018) dataset, and a text classification task for prediction of venues of
scientific papers. Within image classification, we consider two sub-tasks:
Domain Generalization and Sub-population shifts. We defer the synthetic
experiment on IRM, along with the text classification and Sub-population shift
tasks to the Appendix.
For baselines besides IRM, we evaluate the results for standard ERM as well
Information Bottleneck IRM (IB_IRM) (Ahuja et al. 2021). In addition, we
include experiments in the image and language tasks to empirically
characterize the effect of partitioning on ERM and IB_IRM, which we dub as
P-ERM and P-IB_IRM respectively.
An underlying thread for our experiments is the availability to meta-
information that allows us to estimate a notion of distance or similarity
between environments, which P-IRM can then exploit to construct the required
partitions. Specifically, in both housing price prediction and entity
recognition task, our environments are partitioned across time and due to
distribution shifts, we expect environments closer in time to have higher
similarity. Similarly in MetaShift, meta-labels for each image is made
available within the data-set, that allows an explicit notion of the distance
between training and testing environments. In all our experiments we employ
the train-domain validation strategy(Gulrajani and Lopez-Paz 2020) for hyper-
parameter tuning. The code is available at
https://github.com/IbtihalFerwana/pirm and other implementation details are
deferred to Appendix.

###  4.1 Linear Regression

We consider a regression task to predict house prices based on house features
222House Prices Dataset: https://www.kaggle.com/c/house-prices-advanced-
regression-techniques, built across years [1910-2010]. Each data point
consists of 79 predictive features (for instance, number of bedrooms or house
area) and a corresponding target, which is the house price. As pre-processing,
we drop all non-numerical features, samples with missing values and normalize
each feature and price labels to zero mean, unit variance with the samples,
$\\{X_{i},y_{i}\\}_{i}\in(\mathbb{R}^{32}\times\mathbb{R})$.
Experiment Setup To adapt this task to OoD prediction, following (Lin, Zhu,
and Cui 2022), we manually split the training data-set into 10-year segments
and use the house year built as a meta-data for partitioning, with the
intuition being that factors affecting house prices change over time.
For prediction, we consider a linear regression model for the task. Since the
IRM framework $w\circ{\Phi}$ is inherently overparametrized, we fix
$w=1.0\in\mathbb{R}$ and we consider $\Phi\in\mathbb{R}^{32}$ (prediction
($\Phi^{\top}\bm{X}$)) with the Adam optimizer (Kingma and Ba 2015). We
consider 6 training environments corresponding to years [1910-1970], while the
test samples draw from 4 OoD environments [1970-2010]. We expect partitions
closer to the test set to yield better predictors.
Results We report the test MSE error (both average and worst group) over the
set of testing OoD environments, averaged over $5$ random seeds in Table 1. We
find that P-IRM significantly improves the average and worst group OoD error
over IRM. Partitioning also benefits ERM, showing more evidence of a
distribution shift over time, evidence presented in the Appendix. Finally,
note that for the two variants for P-IRM, partitioning performs much better
where we have more samples than parameters.

Model | Training | Avg. MSE | Worst Group MSE
---|---|---|---
ERM | 1910-1970 | 0.475 (0.000) | 1.037 (0.000)
ERM | 1930-1970 | 0.431 (0.000) | 0.963 (0.000)
IRM | 1910-1970 | 0.522 (0.015) | 1.129 (0.038)
P-IRM (partitioned) | 1930-1970 | 0.427 (0.009) | 0.873 (0.024)
P-IRM (conditioned) | 1930-1970 | 0.490 (0.014) | 1.035 (0.034)
Table 1: House Prices Shifts, partitioning demonstrates improvement for both
ERM and IRM, test set is 4 OoD environments consisting of houses built between
1970-2010.

###  4.2 Image Classification

We evaluate P-IRM on a binary image classification task on the MetaShift
dataset (Liang and Zou 2022).
Dataset In MetaShift dataset, each image is associated with a set of tags that
describe the image context (e.g., cat on a rug, cat beside a chair). Thus, for
each given tag (e.g. rug, chair), there is an associated set of images and
these sets can overlap if an image has multiple tags. This structure naturally
induces a graph, with each image context $C_{i}$ denotes a node (or community)
in the graph. This graph is weighted and the weights between nodes is
determined by the number of images that are shared between the communities.
The weights between each pair of communities, $C_{i}$ and $C_{j}$, estimate
the similarity between two communities and are calculated using the
Szymkiewicz-Simpson coefficient, which yields the corresponding adjacency
matrix $\bm{G}$:

$G(i,j)=\tfrac{|C_{i}\cap C_{j}|}{\min(|C_{i}|,|C_{j}|)}$ (6)

Having access to such an undirected weighted graph over sets of images thus
allows us to derive an implicit notion of distance between the corresponding
communities.
Notion of Distance To introduce partitioning, we develop a notion of distance,
which then allows us to quantify the relatedness between training and testing
environments. These environments are assumed to be sets of communities. To
estimate the distance $d$ between any two given nodes/communities, given that
our data is structured as a weighted graph, we can make use of the spectral
embeddings (Belkin and Niyogi 2001). Spectral embeddings are based on graph
Laplacian connectivity (Ng, Jordan, and Weiss 2001). The graph Laplacian
$\bm{L}$ is calculated by $\bm{L}=\bm{D}_{diag}-\bm{G}$, where $\bm{D}_{diag}$
is a diagonal degree matrix of the graph $\bm{G}$. The corresponding
eigenvectors of $L$, $\bm{u}_{1},\dots,\bm{u}_{k}$, computed and normalized to
form the matrix $\bm{U}$, are the corresponding embeddings for the graph. Once
we calculate the spectral embeddings, we measure $d$ between communities as
the euclidean distance between the corresponding spectral embeddings of each
community node. With our notion of distance, we can partition the graph based
on distances between sets of communities and identify subsets of training
communities which are closer to the test environment.
Experiment Setup For all our experiments, we consider the same set of training
communities as in (Liang and Zou 2022), which are split into two environments
in the IRM setting. To introduce partitioning, we assume distances $d$ between
the training environments and the test communities is known/can be estimated
via the meta-labels. For learning the P-IRM model, we consider the training
environment for IRM which is closer to the test set on average, and split it
into two sub-environments. Note that under this split, P-IRM has access to
roughly only half the training samples compared to IRM. To remedy this, we
consider additional data splits wherein we add samples from communities in the
other IRM training environment, that are close to the test set. These
additional samples amount to a percentage $p$ of samples in that environment,
allowing P-IRM access to a slightly larger portion of the training set.
Following (Liang and Zou 2022), we fix the test community to be $dog(shelf)$
and vary distance $d$ between dog train vs test communities. The cat training
set remains unchanged.

| $d=0.17$ | $d=0.54$ | $d=0.81$ | $d=0.92$ | Avg. Performance
---|---|---|---|---|---
ERM | 0.777(0.078) | 0.560(0.179) | 0.493(0.119) | 0.667(0.114) | 0.62425
P-ERM $(p=0)$ | 0.823(0.045) | 0.790(0.086) | 0.387(0.074) | 0.663(0.192) | 0.66575
P-ERM $(p=10)$ | 0.820(0.098) | 0.770(0.057) | 0.493(0.141) | 0.663(0.128) | 0.6865
IRM | 0.757(0.231) | 0.477(0.172) | 0.757(0.110) | 0.687(0.309) | 0.6695
P-IRM $(p=0)$ | 0.960(0.050) | 0.817(0.045) | 0.487(0.083) | 0.650(0.142) | 0.7285
P-IRM $(p=10)$ | 0.710(0.107) | 0.813(0.147) | 0.727(0.087) | 0.690(0.184) | 0.735
IB_IRM | 0.647(0.197) | 0.740(0.171) | 0.750(0.155) | 0.303(0.241) | 0.61
P-IB_IRM $(p=0)$ | 0.663(0.242) | 0.643(0.137) | 0.437(0.289) | 0.617(0.059) | 0.59
P-IB_IRM $(p=10)$ | 0.690(0.340) | 0.790(0.070) | 0.377(0.214) | 0.837(0.160) | 0.6735

Table 2: Domaing Generalization in Metashift. Training environments are $d$
away from testing community $dog(shelf)$, with additional samples up to
percentage $p\in\\{0,10\\}$ for partitioned models. Results for $p=25$ are in
Table 5 (in Appendix)

.

Results For all experiments, we report the binary classification accuracy
averaged over 3 seeds, with the randomness solely arising from the learning
algorithm. We compare the performance of P-IRM against IRM, as well other
benchmarks and their corresponding partitioned versions in table 2. In most
experiments, especially those with higher deviation between the training and
testing data, models with partitioning tend to perform better.

###  4.3 Named Entity Recognition (NER)

Distributional shifts are common in language tasks, given that societal
changes are known to influence language usage over time. These changes are
also reflected in word embeddings (words vectors to represent language) (Garg
et al. 2018). Within this context, we explore effects of partitioning
(Lazaridou et al. 2021; Luu et al. 2021).
Experiment Setup We consider the SciERC (Luan et al. 2018) dataset, which
consists of CS publications from 1980 to 2016. The specific task is Named
Entity Recognition, a multi-class classification task, that labels each
scientific mention in a sentence into six possible categories (Task, Method,
Evaluation Metric, Material, Other-Scientific-Term, or Generic).The training
set comprises of years from 1980-2009 and we test the model on data obtained
between 2010-2016, with an intention to study distribution shift over time.
For creating the training environments, we split training years into smaller
intervals, 1990-2009, 2000-2009 and 2005-2009, such that each interval has
roughly the same number of samples. For partitioning, we consider contiguous
partitions of time intervals, based on the intuition that vocabularies in text
have higher overlap when closer in time (Gururangan et al. 2020). For building
the model, we train a classifier over the BERT pretrained language model
(Devlin et al. 2019). Due to high sample complexity, we also consider the
conditioned P-IRM method that makes use of all training environments.
Results We report the classification accuracy, averaged over $3$ seeds in
table 3. We find that both variants of P-IRM indeed improve performance over
IRM. Additionally, we find that leveraging more training data using
conditioned P-IRM leads to marginally better predictors, when compared against
standard partitioning. Comparisons against IB_IRM as well as ERM demonstrate
that partitioning can improve efficacy of other learning algorithms as well.

Model | # envs | Training | Accuracy (2010-2016)
---|---|---|---
ERM | 4 | 1980-2009 | 0.800 (0.012)
P-ERM | 3 | 1990-2009 | 0.804 (0.020)
P-ERM | 2 | 2000-2009 | 0.804 (0.016)
IRM | 4 | 1980-2009 | 0.795 (0.005)
P-IRM (partitioned) | 3 | 1990-2009 | 0.795 (0.017)
P-IRM (partitioned) | 2 | 2000-2009 | 0.807 (0.005)
P-IRM (conditioned) | 3 | 1990-2009 | 0.812 (0.008)
P-IRM (conditioned) | 2 | 2000-2009 | 0.807 (0.015)
IB_IRM | 4 | 1980-2009 | 0.800 (0.010)
P-IB_IRM (partitioned) | 3 | 1990-2009 | 0.800 (0.015)
P-IB_IRM (partitioned) | 2 | 2000-2009 | 0.794 (0.015)
P-IB_IRM (conditioned) | 3 | 1990-2009 | 0.807 (0.008)
P-IB_IRM (conditioned) | 2 | 2000-2009 | 0.805 (0.020)
Table 3: Language Shifts in SciERC dataset. Partitioning improves performance,
with (1990-2009) consistently optimal across all learning algorithms.

##  5 Discussion

In this work, we propose partial invariance, as a relaxation of IRM objective,
which allows us to explore a subtle trade-off in invariant models, namely
accessing more domains at the cost of a smaller permissible invariant feature
set. We then verify, with experiments across multiple domains, that when
feasible, partitioning can indeed improve upon IRM as well as other learning
frameworks.
We note that the proposed framework is naturally limited by the available
information about training/deployment domains. While distribution shifts
across time allows for partitions to be contiguous time intervals, finding
appropriate partitions is non-trivial under complex shift topologies. In that
sense, our work is the first step towards understanding the need for training
domain selection in invariant learning. Thus, developing general heuristics
for identifying the right partition is an important direction of future work.
Second, we note that the conditional variant of P-IRM provides tangible gains
in low data regimes, and it is of interest to study the nature of the
accessible feature set as well as the associated sample complexities.

## References

  * Adragna et al. (2020) Adragna, R.; Creager, E.; Madras, D.; and Zemel, R. S. 2020.  Fairness and Robustness in Invariant Learning: A Case Study in Toxicity Classification.  _CoRR_ , abs/2011.06485.
  * Ahuja et al. (2021) Ahuja, K.; Caballero, E.; Zhang, D.; Gagnon-Audet, J.-C.; Bengio, Y.; Mitliagkas, I.; and Rish, I. 2021.  Invariance Principle Meets Information Bottleneck for Out-of-Distribution Generalization.  In _Advances in Neural Information Processing Systems_.
  * Ahuja et al. (2020a) Ahuja, K.; Shanmugam, K.; Varshney, K. R.; and Dhurandhar, A. 2020a.  Invariant Risk Minimization Games.  In _Proceedings of the 37th International Conference on Machine Learning (ICML’20)_ , volume 119, 145–155. PMLR.
  * Ahuja et al. (2020b) Ahuja, K.; Wang, J.; Dhurandhar, A.; Shanmugam, K.; and Varshney, K. R. 2020b.  Empirical or Invariant Risk Minimization? A Sample Complexity Perspective.  In _Proceeding of the 8th International Conference on Learning Representations (ICLR’20)_.
  * Arjovsky et al. (2019) Arjovsky, M.; Bottou, L.; Gulrajani, I.; and Lopez-Paz, D. 2019.  Invariant Risk Minimization.  arXiv:1907.02893 [stat.ML].
  * Beery, V. Horn, and Perona (2018) Beery, S.; V. Horn, G.; and Perona, P. 2018.  Recognition in Terra Incognita.  In _Proceedings of the European Conference on Computer Vision (ECCV)_ , 456–473.
  * Belkin and Niyogi (2001) Belkin, M.; and Niyogi, P. 2001.  Laplacian eigenmaps and spectral techniques for embedding and clustering.  _Advances in neural information processing systems_ , 14.
  * Bellot and van der Schaar (2020) Bellot, A.; and van der Schaar, M. 2020.  Accounting for Unobserved Confounding in Domain Generalization.  arXiv:2007.10653 [stat.ML].
  * Ben-David et al. (2010) Ben-David, S.; Blitzer, J.; Crammer, K.; Kulesza, A.; Pereira, F.; and Vaughan, J. 2010.  A theory of learning from different domains.  _Machine Learning_ , 79: 151–175.
  * Ben-Tal, El Ghaoui, and Nemirovski (2009) Ben-Tal, A.; El Ghaoui, L.; and Nemirovski, A. 2009.  _Robust Optimization_.  Princeton Series in Applied Mathematics. Princeton University Press.
  * Devlin et al. (2019) Devlin, J.; Chang, M.-W.; Lee, K.; and Toutanova, K. 2019.  BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , 4171–4186.
  * Duchi, Glynn, and Namkoong (2021) Duchi, J.; Glynn, P.; and Namkoong, H. 2021.  Statistics of Robust Optimization: A Generalized Empirical Likelihood Approach.  _Mathematics of Operations Research_ , 46(3).
  * Gallacher (2021) Gallacher, J. D. 2021.  Leveraging cross-platform data to improve automated hate speech detection.  arXiv:2102.04895 [CS.CL].
  * Garg et al. (2018) Garg, N.; Schiebinger, L.; Jurafsky, D.; and Zou, J. 2018.  Word embeddings quantify 100 years of gender and ethnic stereotypes.  _Proceedings of the National Academy of Sciences_ , 115(16): E3635–E3644.
  * Gong et al. (2016) Gong, M.; Zhang, K.; Liu, T.; Tao, D.; Glymour, C.; and Schölkopf, B. 2016.  Domain Adaptation with Conditional Transferable Components.  In _Proceedings of the 33rd International Conference on Machine Learning (ICML’16)_ , volume 48, 2839–2848.
  * Gulrajani and Lopez-Paz (2020) Gulrajani, I.; and Lopez-Paz, D. 2020.  In Search of Lost Domain Generalization.  In _Proceeding of the 8th International Conference on Learning Representations (ICLR’20)_.
  * Gulrajani and Lopez-Paz (2020) Gulrajani, I.; and Lopez-Paz, D. 2020.  In Search of Lost Domain Generalization.  _CoRR_ , abs/2007.01434.
  * Gururangan et al. (2020) Gururangan, S.; Marasović, A.; Swayamdipta, S.; Lo, K.; Beltagy, I.; Downey, D.; and Smith, N. A. 2020.  Don’t Stop Pretraining: Adapt Language Models to Domains and Tasks.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , 8342–8360.
  * He et al. (2016) He, K.; Zhang, X.; Ren, S.; and Sun, J. 2016.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 770–778.
  * Heinze-Deml, Peters, and Meinshausen (2018) Heinze-Deml, C.; Peters, J.; and Meinshausen, N. 2018.  Invariant Causal Prediction for Nonlinear Models.  _Journal of Causal Inference_ , 6(2).
  * Jin, Barzilay, and Jaakkola (2020) Jin, W.; Barzilay, R.; and Jaakkola, T. S. 2020.  Domain Extrapolation via Regret Minimization.  _CoRR_ , abs/2006.03908.
  * Kamath Pritish and Srebro (2021) Kamath Pritish, D. S., Akilesh Tangella; and Srebro, N. 2021.  Does Invariant Risk Minimization Capture Invariance?  In _Proceedigns of the International Conference on Artificial Intelligence and Statistics_ , 4069–4077. PMLR.
  * Kingma and Ba (2015) Kingma, D. P.; and Ba, J. 2015.  Adam: A Method for Stochastic Optimization.  In Bengio, Y.; and LeCun, Y., eds., _3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings_.
  * Koh et al. (2021) Koh, P. W.; Sagawa, S.; Marklund, H.; Xie, S. M.; Zhang, M.; Balsubramani, A.; Hu, W.; Yasunaga, M.; Phillips, R. L.; Gao, I.; Lee, T.; David, E.; Stavness, I.; Guo, W.; Earnshaw, B.; Haque, I.; Beery, S. M.; Leskovec, J.; Kundaje, A.; Pierson, E.; Levine, S.; Finn, C.; and Liang, P. 2021.  WILDS: A Benchmark of in-the-Wild Distribution Shifts.  In Meila, M.; and Zhang, T., eds., _Proceedings of the 38th International Conference on Machine Learning_ , volume 139 of _Proceedings of Machine Learning Research_ , 5637–5664. PMLR.
  * Krueger et al. (2021) Krueger, D.; Caballero, E.; Jacobsen, J.-H.; Zhang, A.; Binas, J.; Zhang, D.; Le Priol, R.; and Courville, A. 2021.  Out-of-Distribution Generalization via Risk Extrapolation.  In _Proceedings of the 38th International Conference on Machine Learning (ICML’21)_ , 5815–5826. PMLR.
  * Lake et al. (2017) Lake, B. M.; Ullman, T. D.; Tenenbaum, J. B.; and Gershman, S. J. 2017.  Building machines that learn and think like people.  _Behavioral and Brain Sciences_ , 40: e253.
  * Lazaridou et al. (2021) Lazaridou, A.; Kuncoro, A.; Gribovskaya, E.; Agrawal, D.; Liska, A.; Terzi, T.; Gimenez, M.; de Masson d’Autume, C.; Ruder, S.; Yogatama, D.; Cao, K.; Kociský, T.; Young, S.; and Blunsom, P. 2021.  Pitfalls of Static Language Modelling.  arXiv:2102.01951 [cs.CL].
  * Lee and Raginsky (2018) Lee, J.; and Raginsky, M. 2018.  Minimax Statistical Learning with Wasserstein Distances.  In _Proceedings of the 32nd International Conference on Neural Information Processing Systems (NIPS’18)_ , 2692–2701.
  * Li et al. (2018) Li, Y.; Gong, M.; Tian, X.; Liu, T.; and Tao, D. 2018.  Domain Generalization via Conditional Invariant Representation.  _Proceedings of the 32nd Association for the Advancement of Artificial Intelligence (AAAI’18)_ , 31(1).
  * Liang and Zou (2022) Liang, W.; and Zou, J. 2022.  Metashift: A dataset of datasets for evaluating contextual distribution shifts and training conflicts.  In _International Conference on Learning Representations, ICLR 2022_.
  * Lin, Zhu, and Cui (2022) Lin, Y.; Zhu, S.; and Cui, P. 2022.  ZIN: When and How to Learn Invariance by Environment Inference?  _arXiv preprint arXiv:2203.05818_.
  * Luan et al. (2018) Luan, Y.; He, L.; Ostendorf, M.; and Hajishirzi, H. 2018.  Multi-Task Identification of Entities, Relations, and Coreferencefor Scientific Knowledge Graph Construction.  In _Proc. Conf. Empirical Methods Natural Language Process. (EMNLP)_.
  * Luu et al. (2021) Luu, K.; Khashabi, D.; Gururangan, S.; Mandyam, K.; and Smith, N. A. 2021.  Time Waits for No One! Analysis and Challenges of Temporal Misalignment.  ArXiv preprint arXiv:2111.07408.
  * Mahajan, Tople, and Sharma (2021) Mahajan, D.; Tople, S.; and Sharma, A. 2021.  Domain Generalization using Causal Matching.  In _Proceedings of the 38th International Conference on Machine Learning (ICML’21)_ , volume 139, 7313–7324. PMLR.
  * Mani, Varshney, and Pentland (2021) Mani, A.; Varshney, L. R.; and Pentland, A. 2021.  Quantization Games on Social Networks and Language Evolution.  _IEEE Transactions on Signal Processing_ , 69: 3922–3934.
  * Marcus (2018) Marcus, G. 2018.  Deep Learning: A Critical Appraisal.  arXiv:1801.00631 [CS.AI].
  * Mohri, Sivek, and Suresh (2019) Mohri, M.; Sivek, G.; and Suresh, A. T. 2019.  Agnostic Federated Learning.  In _Proceedings of the 36th International Conference on Machine Learning (ICML’19)_ , volume 97, 4615–4625. PMLR.
  * Neal (1996) Neal, R. M. 1996.  _Bayesian Learning for Neural Networks_.  Berlin, Heidelberg: Springer-Verlag.  ISBN 0387947248.
  * Ng, Jordan, and Weiss (2001) Ng, A.; Jordan, M.; and Weiss, Y. 2001.  On spectral clustering: Analysis and an Algorithm.  _Advances in Neural Information Processing Systems_ , 14.
  * Pearl (2009) Pearl, J. 2009.  Causal inference in statistics: An overview.  _Statistics Surveys_ , 3(none): 96 – 146.
  * Peters, Bühlmann, and Meinshausen (2016) Peters, J.; Bühlmann, P.; and Meinshausen, N. 2016.  Causal inference by using invariant prediction: identification and confidence intervals.  _Journal of the Royal Statistical Society. Series B (Statistical Methodology)_ , 78(5): 947–1012.
  * Peters, Bühlmann, and Meinshausen (2015) Peters, J.; Bühlmann, P.; and Meinshausen, N. 2015.  Causal inference using invariant prediction: identification and confidence intervals.  Preprint.
  * Peyrard et al. (2021) Peyrard, M.; Ghotra, S. S.; Josifoski, M.; Agarwal, V.; Patra, B.; Carignan, D.; Kiciman, E.; and West, R. 2021.  Invariant Language Modeling.  _CoRR_ , abs/2110.08413.
  * Radford et al. (2019) Radford, A.; Wu, J.; Child, R.; Luan, D.; Amodei, D.; Sutskever, I.; et al. 2019\.  Language models are unsupervised multitask learners.  _OpenAI blog_ , 1(8): 9.
  * Rosenfeld, Ravikumar, and Risteski (2020) Rosenfeld, E.; Ravikumar, P. K.; and Risteski, A. 2020.  The Risks of Invariant Risk Minimization.  In _Proceeding of the 8th International Conference on Learning Representations (ICLR’20)_.
  * Sanh et al. (2019) Sanh, V.; Debut, L.; Chaumond, J.; and Wolf, T. 2019.  DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter.
  * Schölkopf et al. (2012) Schölkopf, B.; Janzing, D.; Peters, J.; Sgouritsa, E.; Zhang, K.; and Mooij, J. 2012.  On causal and Anticausal Learning.  In _Proceedings of the 29th International Conference on Machine Learning (ICML’12)_ , 1255–1262.
  * Shui, Wang, and Gagné (2021) Shui, C.; Wang, B.; and Gagné, C. 2021.  On the benefits of representation regularization in invariance based domain generalization.  _CoRR_ , abs/2105.14529.
  * Vapnik (2013) Vapnik, V. 2013.  _The Nature of Statistical Learning Theory_.  Springer Science and Business Media.
  * Xie et al. (2021) Xie, C.; Ye, H.; Chen, F.; Liu, Y.; Sun, R.; and Li, Z. 2021.  Risk Variance Penalization.  arXiv:2006.07544 [cs.LG].
  * Yu et al. (2022) Yu, R.; Zhu, H.; Li, K.; Hong, L.; Zhang, R.; Ye, N.; Huang, S.-L.; and He, X. 2022\.  Regularization Penalty Optimization for Addressing Data Quality Variance in OoD Algorithms.  _Proceedings of the AAAI Conference on Artificial Intelligence_ , 36(8): 8945–8953.
  * Zhang, Gong, and Schoelkopf (2015) Zhang, K.; Gong, M.; and Schoelkopf, B. 2015.  Multi-Source Domain Adaptation: A Causal View.  In _Proceedings of the 29th Association for the Advancement of Artificial Intelligence Conference (AAAI’15)_.

##  Appendix A Proofs

###  A.1 Proof of Lemma 1

The proof follows as a consequence of the parametrization. Under the MSE loss,
note that the expected risk takes the following form:

$\displaystyle{R}^{e}(w\circ\Phi)=\bm{E}_{\bm{y}^{e},\bm{X}^{e}}(\bm{y}^{e}-w.\Phi^{\top}\bm{X}^{e})^{2}\mbox{,}$
(7)

wherein the scalar $w$ is fixed at 1.0. The gradient penalty w.r.t. $w$ in
environment $e$ can then be obtained as:

$\displaystyle\begin{split}&\
\|\nabla_{w,w=1.0}{R}^{e^{\prime}}(w\circ\Phi)\|=|\nabla_{w,w=1.0}\bm{E}_{\bm{y}^{e},\bm{X}^{e}}(\bm{y}^{e}-w.\Phi^{\top}\bm{X}^{e})^{2}|\\\
&=|\bm{E}_{\bm{y}^{e},\bm{X}^{e}}[\nabla_{w,w=1.0}(\bm{y}^{e}-w.\Phi^{\top}\bm{X}^{e})]^{2}|\\\
&=|\bm{E}_{\bm{y}^{e},\bm{X}^{e}}[2(w.\Phi^{\top}\bm{X}^{e}-\bm{y}^{e})\Phi^{\top}\bm{X}^{e}]|\\\
&=|\bm{E}_{{\epsilon_{y}}^{e},\bm{X}^{e}}[2(\Phi-W^{e})^{\top}\bm{X}^{e}-{\epsilon_{y}}^{e}){\bm{X}^{e}}^{\top}\Phi]|\\\
&=|\bm{E}_{\bm{X}^{e}}[2(\Phi-W^{e})^{\top}(\bm{X}^{e}{\bm{X}^{e}}^{\top})\Phi]+0|\\\
&=|2(\Phi-W^{e})^{\top}\bm{E}_{\bm{X}^{e}}[(\bm{X}^{e}{\bm{X}^{e}}^{\top})]\Phi|\\\
&=|2(\Phi-W^{e})^{\top}\Phi|\\\
&=2|\sum_{i=1}^{c}(\Phi_{i}^{2}-w^{e}_{i}.\Phi_{i})|\mbox{,}\\\ \end{split}$
(8)

wherein the simplifications follow through due to feature independence,
normalization and zero mean noise assumption. Notice that as per the
constraints in (2), we need the risk penalty term above to be equal to zero
for all training environments:

$|\sum_{i=1}^{c}(\Phi_{i}^{2}-w^{e}_{i}.\Phi_{i})|=0\ \forall\
e\in\mathcal{E}_{tr}.$

Then note from a risk minimization incentive, we naturally have
$\Phi_{1}=w_{inv}$ without incurring any penalty. Thus the penalty boils to:

$|\sum_{i=2}^{c}(\Phi_{i}^{2}-w^{e}_{i}.\Phi_{i})|=0\ \forall\
e\in\mathcal{E}_{tr}.$

But from Assumption 1, we have an environment $e$ in the training set in which
IRM is sufficient, i.e. $w^{e}_{i}=0\ \forall\ i\neq 1$. Thus, the constraint
in that environment equates to:

$|\sum_{i=2}^{c}(\Phi_{i}^{2}-0.\Phi_{i})|=|\sum_{i=2}^{c}\Phi_{i}^{2}|=0\
\forall\ e\in\mathcal{E}_{tr}.$

Thus to satisfy this constraint, $\Phi_{i}=0\ \forall\ i\neq 1$, which means
the set of feasible solutions is comprised solely of the perfectly invariant
predictor, as claimed.

###  A.2 Proof of Lemma 2

Under the uniform feature model, we sample $\mathcal{E}_{tr}$ such that
$|\mathcal{E}_{tr}|=t$. Notice that the cardinality $m$ of the required
partition $\mathcal{E}_{partition}:=\\{e\in\mathcal{E}_{tr}|\omega(e)=1\\}$ is
random variable.
We note that for an environment sample $e\in\mathcal{E}_{tr}$, for a given
reference environment $e^{ref}$, we have from our analysis in the previous
theorem that $P(\|W^{e}-W^{e_{ref}}\|\leq\delta)=P(E_{1})+P(E_{2})$ and that
$P(E_{1})>pP(E_{2})$, wherein $p\geq\frac{(c-1-\delta)\alpha}{\delta}>>1$. Let
$n=c-2$ for brevity. Then, we have the following approximation:

$\displaystyle P(\|W^{e}-W^{e_{ref}}\|\leq\delta)\approx P(E_{1})$
$\displaystyle\geq\gamma=\frac{1}{k\sqrt{2n}}\exp(-nD(\delta/n\|1/\alpha k)),$

where the final result follows from standard anti-concentration bounds on a
Binomial distribution. Then it is easy to see that given
$|\mathcal{E}_{tr}|=t$, $P(|\mathcal{E}_{partition}|\geq
m)=P(|\sum_{e\in\mathcal{E}_{tr}}\bm{1}[\omega(e)=1]\geq m)]\geq
P(\sum_{j=1}^{t}D_{j}\geq m]$ where $D_{j}$ is Bernoulli random variable
distributed as $D_{j}\sim\text{Bern}(\gamma)$.

On the RHS, we get a sum of i.i.d. Bernoulli variables for and is thus,
$\sum_{j=1}^{t}D_{j}\sim\text{Bin}(n,\gamma)$. Let $M=\sum_{j=1}^{t}D_{j}$ and
notice that we need that $M>m$ with high probability. To achieve this, we
first upper bound the probability of event $M<m$ using Chernoff’s tail bound
and derive conditions under which this upper bound is small. Specifically:

$P(M<m)<\text{exp}(-tD\bigg{(}\frac{m}{t}\|\gamma\bigg{)})<\epsilon,$ (9)

where $D(a\|b)=a\log(\frac{a}{b})+(1-a)\log(\frac{1-a}{1-b})$. Next, assume
that $t=Cm$. Then we can simplify the inequality:

$\displaystyle-
Cm\bigg{(}\frac{1}{C}\log\frac{\frac{1}{\gamma}}{C}+\frac{C-1}{C}\log\frac{(C-1)\frac{1}{\gamma}}{C(\frac{1}{\gamma}-1)}\bigg{)}<\log(\epsilon)$
(10a)
$\displaystyle=>\bigg{(}\log\frac{\frac{1}{\gamma}}{C}+({C-1})\log\frac{(C-1)\frac{1}{\gamma}}{C(\frac{1}{\gamma}-1)}\bigg{)}>1/m\log(1/\epsilon)$
(10b)
$\displaystyle=>\bigg{(}C\log\frac{(C-1)\frac{1}{\gamma}}{C(\frac{1}{\gamma}-1)}-\log\frac{C-1}{\frac{1}{\gamma}-1}\bigg{)}>1/m\log(1/\epsilon)$
(10c)
$\displaystyle\asymp\bigg{(}C\log\frac{(C-1)\frac{1}{\gamma}}{C(\frac{1}{\gamma}-1)}\bigg{)}>1/m\log(1/\epsilon).$
(10d)

Our inequality will be satisfied if a) $C>\frac{1}{\gamma}\log(1/\epsilon)$
and b) $\log(\frac{(C-1)\frac{1}{\gamma}}{C(\frac{1}{\gamma}-1)})>\gamma m$.
Then we can show that for sufficiently large $\frac{1}{\gamma}m$, condition b)
roughly amounts $C>c_{i}(1+\frac{1}{m-1})$. So if
$C>\max\\{\frac{1}{\gamma}\log(1/\epsilon),\frac{1}{\gamma}(1+\frac{1}{m-1}\\}\sim\frac{1}{\gamma}\log(1/\epsilon)$
(for small $\epsilon$), then we have $P(M<m)<\epsilon$. Hence, for
$t=Cm\sim\frac{1}{\gamma}m\log(1/\epsilon),$ we get that
$M=\sum_{j=1}^{t}D_{j}$ with high probability. But note that we already have
$P(|\mathcal{E}_{partition}|\geq m)\geq P(\sum_{j=1}^{t}D_{j}\geq m]$. Thus,
$P(|\mathcal{E}_{partition}|\geq m)>1-\epsilon$.

###  A.3 Proof of Theorem 1

We begin by restating the theorem.

######  Theorem.

$\bm{1}$ Under the model (4), under Assumption 2, with access to oracle
$\omega(e)=\mathbf{1}[\|W^{e_{ref}}-W^{e}\|_{0}\leq\delta]$ and
$\delta<(c-2)/2$, isolate
$\mathcal{E}_{partition}:=\\{e\in\mathcal{E}_{tr}|\omega(e)=1\\}\cup\\{e_{ref}\\}\subseteq\mathcal{E}_{tr}$.
Next, let $|A_{i}|=k$, where $A_{i}$ is the set corresponding to the feature
weight $w_{i}^{e_{ref}}$ of interest. Then, if the sets $\\{A_{j}\\}\forall\
j\in\\{2,\ldots,c\\}\setminus\ \\{i\\}$ satisfy $|A_{j}|>\alpha k$ for some
$\alpha>1$, we have with probability greater or equal to
$(\frac{p}{p+1})^{|\mathcal{E}_{partition}|}$, where
$p\geq\frac{(c-1-\delta)\alpha}{\delta}$, the IRM solution over set
$\mathcal{E}_{partition}$ will recover the feature of interest
$w_{i}^{e_{ref}}$.

The proof sketch is as follows. We first characterize the probability of error
in using the oracle as an indicator for the partition membership. Assuming the
partition is identified, we can then directly apply Lemma 1 to obtain the
desired result. We begin by characterizing the two possible cases that arise
when oracle $\omega(e)$ = 1.

  * •

$E_{1}$: If feature at level $i$ remains unchanged, this means that out of
$c-2$ features (discarding the feature at level $i$ and the invariant feature
at level $1$), a maximum of up-to $\delta$ features changed in the worst case.

  * •

$E_{2}$: If feature at level $i$ changed, this means that out of $c-2$
features (discarding the feature at level $i$ and the invariant feature at
level $1$), a maximum of up-to $\delta-1$ features changed in the worst case.

Note that both of the events can be modelled as a sum of Bernoulli random
variables with different probabilities of success. Consider
$B_{j}\sim\text{Bern}(1-1/|A_{j}|)$, which indicates whether the feature value
at level $j$ changed under the uniform model. In the first case, we model the
conditional probability as:

$\displaystyle\begin{split}&\
P(E_{1})=P(w^{e}_{i}=w_{i}^{e_{ref}}|\omega(e)=1)\\\
&=(1/k)P(1\leq\sum_{{j=2},j\neq i}^{c}B_{j}\leq\delta),\end{split}$ (11)

while in the alternate case,

$\displaystyle\begin{split}&P(E_{2})=P(w^{e}_{i}\neq
w_{i}^{e_{ref}}|\omega(e)=1)\\\ &=(1-1/k)P(0\leq\sum_{{j=2},j\neq
i}^{c}B_{j}\leq\delta-1).\end{split}$ (12)

Next, we make use of the assumption on cardinalities and analyze the
probabilities corresponding to each term. We note that:

$P(E_{1})=\sum_{m=1}^{\delta}(1/k)P(\sum_{{j=2},j\neq i}^{c}B_{j}=m).$

Similarly:

$P(E_{2})=\sum_{m=0}^{\delta-1}(1-1/k)P(\sum_{{j=2},j\neq i}^{c}B_{j}=m)$

For brevity, we represent each individual term on the right hand side of the
summations as $P_{a,m}$ $(a=1\text{ or }2\text{ depending on the event})$,
corresponding to that value of $m$. Next, we analyze the ratio for the final
two terms in either sequence. Let $c-2=n$ and note that using the assumption
on cardinalities of the feature sets, we can bound the ratio as follows:

$\frac{P_{1,\delta}}{P_{2,{\delta-1}}}\geq\frac{(1/k)({\binom{n}{\delta}}(1-1/\alpha
k)^{\delta}(1/\alpha k)^{n-\delta})}{(1-1/k)({\binom{n}{\delta-1}}(1-1/\alpha
k)^{\delta-1}(1/\alpha k)^{n+1-\delta})},$

wherein the greater than equal to sign holds since cardinality of each feature
set is greater than or equal to $\alpha k$. Simplifying the terms, we get:

$\frac{P_{1,\delta}}{P_{2,{\delta-1}}}\geq\frac{({n+1-\delta)}(\alpha
k-1)}{(k-1)({\delta)}}\geq\frac{(c-1-\delta)\alpha}{\delta}=p.$

Note that this ratio increases as $\delta$ reduces and therefore, we conclude
that:

$P(E_{1})/P(E_{2})=\frac{\sum_{m=1}^{\delta}P_{1,m}}{\sum_{m=1}^{\delta-1}P_{2,m}}\geq
p.$

Thus, we obtain that $P_{error}=P(E_{2})\leq 1/(p+1)$. Thus, for our training
partition $\mathcal{E}_{partition}=\\{e\in\mathcal{E}_{tr}|\omega(e)=1\\}$,
with probability greater than $(\frac{p}{p+1})^{|\mathcal{E}_{partition}|}$,
we will learn an accurate partition. Since within this partition the feature
weight corresponding to $\bm{x}_{i}$ is same as the reference for all
environments, we obtain the required result from Lemma 1.

###  A.4 Proof of Theorem 2

First, we restate the theorem more formally.

######  Theorem.

$\bm{2}$ (Formal) Assume we observe samples $(\tilde{\bm{X}}^{e},\bm{y}^{e})$
as per equation (5), with environments $e\in\mathcal{E}_{tr}$ sampled as per
equation (4) and let
$\mathcal{E}_{partition}:=\\{e\in\mathcal{E}_{tr}|\omega(e)=1\\}$. Let
$\bm{\Phi}\in\mathbb{R}^{d\times d}$ have rank $r$. Then sampling
$|\mathcal{E}_{tr}|>\frac{1}{\gamma}(d-r+d/r)\log(1/\epsilon)$ ensures that
partition cardinality $|\mathcal{E}_{partition}|>d-r+d/r$ with probability ¿
$1-\epsilon$. Furthermore, if $e\in\mathcal{E}_{tr}$ lie in linear general
position of degree $r$ (Assumption 3), then with probability greater than or
equal to $(\frac{p}{p+1})^{|\mathcal{E}_{partition}|}$, where
$p\geq\frac{(c-1-\delta)\alpha}{\delta}$, the oracle identifies
$\mathcal{E}_{partition}$ such that we have:

$\Phi\bm{E}_{\tilde{\bm{X}}^{e}}[\tilde{\bm{X}}^{e}(\tilde{\bm{X}}^{e})^{\top}]\Phi^{\top}W=\Phi\bm{E}_{\tilde{\bm{X}}^{e},\bm{y}^{e}}[\tilde{\bm{X}}^{e}\bm{y}^{e}],$

holds for all $e\in\mathcal{E}_{partition}$ iff $\bm{\Phi}$ elicits an
invariant predictor $\bm{\Phi}^{\top}W$ $\forall\ e\in\mathcal{E}_{all}$ whose
feature weights satisfy $w^{e}_{i}=w^{e_{ref}}_{i}$.

We begin by recollecting the requisite tools from (Arjovsky et al. 2019).

######  Assumption 3.

With observation model as in equation (5), a set of training environments
$\mathcal{E}_{partition}\subseteq\mathcal{E}_{tr}$ lie in linear general
position of degree $r$ if $|\mathcal{E}_{partition}|>d-r+d/r$ for some
$r\in\mathbb{N}$, $r<d$, and for all non-zero $X\in\mathbb{R}^{d}$:

$\displaystyle\text{dim}\bigg{(}\text{span}\bigg{(}\big{\\{}\bm{E}_{\tilde{\bm{X}}^{e}}[\tilde{\bm{X}}^{e}(\tilde{\bm{X}}^{e})^{\top}]X-\bm{E}_{\tilde{\bm{X}}^{e},\tilde{\epsilon}_{y}}[\tilde{\bm{X}}^{e}\tilde{\epsilon}_{y}]\big{\\}}_{e\in\mathcal{E}_{partition}}\bigg{)}\bigg{)}$
$\displaystyle\quad>d-r.$

Intuitively, this assumption states that we require the training environments
in our partition partition $\mathcal{E}_{tr}$ to be sufficiently diverse, with
limited co-linearity.

Next, recall that in our setup, instead of directly observing the training
partition for a given level $i$ and feature weight $v$, we need to identify
the a subset $\mathcal{E}_{partition}$ from the available set of training
environments $\mathcal{E}_{tr}$. Thus, for the training partition to be at
least of size $m$, we need certain conditions on $|\mathcal{E}_{tr}|$. We
characterize this in the following result.

######  Lemma 2.

Under environment sampling as per equation (4), if the cardinality of observed
environments $|\mathcal{E}_{tr}|=n\sim\frac{1}{\gamma}m\log(1/\epsilon)$, for
the subset that satisfies the oracle distance condition i.e.
$\mathcal{E}_{partition}:=\\{e\in\mathcal{E}_{tr}|\omega(e)=1\\}$, we have
that $P(|\mathcal{E}_{partition}|\geq m)>1-\epsilon$.

Having obtained the correct partitioning with high probability, we call upon
the out of domain generalization result from (Arjovsky et al. 2019).

######  Proposition 1 (Theorem 9 in (Arjovsky et al. 2019)).

Assume that

$\begin{gathered}Y^{e}=(Z_{1}^{e})^{\top}\beta+\epsilon^{e},\epsilon^{e}\perp
Z_{1}^{e},\ \bm{E}[{\epsilon}^{e}]=0\\\
X^{e}=S(Z_{1}^{e},Z_{2}^{e}).\end{gathered}$ (13)

Here, $\beta\in\mathbb{R}^{c}$, $Z_{1}^{e}$ takes values in $\mathbb{R}^{c}$,
$Z_{2}^{e}$ takes values in $\mathbb{R}^{1}$ and
$S\in\mathbb{R}^{d\times(c+q)}$. Assume that the $Z_{1}$ component of $S$ is
invertible: that there exists $\tilde{S}\in\mathbb{R}^{c\times d}$ such that
$\tilde{S}(S(z_{1},z_{2}))=z_{1}$, for all $z_{1}\in\mathbb{R}^{c}$,
$z_{2}\in\mathbb{R}^{q}$. Let $\Phi\in\mathbb{R}^{d\times d}$ have rank $r$.
Then, if atleast $d-r+d/r$ training environments
$\mathcal{E}_{tr}\subseteq\mathcal{E}_{all}$ lie in linear general position of
degree $r$, then we have:

$\Phi\bm{E}_{{X}^{e}}[{X}^{e}({X}^{e})^{\top}]\Phi^{\top}w=\Phi\bm{E}_{{X}^{e},Y^{e}}[{X}^{e}Y^{e}],$

holds for $\forall e\in\mathcal{E}_{tr}$ iff $\Phi$ elicits an invariant
predictor $\Phi^{\top}w$ for all $\ e\in\mathcal{E}_{all}$.

We now provide the proof of Theorem 2. First, let $m=d-r+d/r$. Then from Lemma
2, we know that sampling $n\sim\frac{1}{\gamma}m\log(1/\delta)$ environments
gives us $|\mathcal{E}_{partition}|\geq m$ with high probability. From
Assumption $3$, we also have that environments in $\mathcal{E}_{partition}$
lie in linear general position of degree $r$. Finally, note that our scheme is
contingent on the oracle correctly identifying the required partition, which
happens with probability greater than
$(\frac{p}{p+1})^{|\mathcal{E}_{partition}|}$, as noted in Theorem 1. Armed
with these, we can apply Proposition 1 to our learning setup in equation (5)
to learn an predictor $\bm{\Phi},W$ that can recover features
${\bm{X}}_{inv}^{e}$ and corresponding desired weights $W_{inv}^{e}$ $\forall
e\in\mathcal{E}_{all}$ which satisfy $w^{e}_{i}=w^{e_{ref}}_{i}$.

##  Appendix B Additional Experiments

In all our experiments when implementing IRM/P-IRM, we keep the penalty
parameter sufficiently high $\lambda=10^{2}/10^{3}$. The rationale for this is
to have $\lambda$ high enough so that the invariance penalty term dominates
the fidelity loss term and features are close to invariant.

###  B.1 Synthetic Experiment

The experimental setting is adapted from (Arjovsky et al. 2019). We assume the
following generative model:

$\begin{gathered}\bm{X}_{1}\leftarrow N(\bm{0},e^{2}I),\bm{X}_{2}\leftarrow
N(\bm{0},e^{2}I),\\\ {c(e)\in\\{0,1\\},\epsilon\sim N(\bm{0},e^{2}I)}\\\
\bm{y}\leftarrow\bm{X}_{1}^{T}W_{1}+\bm{X}_{2}^{T}(c(e)W_{2})+\epsilon,\
\epsilon\perp\bm{X}_{1},\bm{X}_{2}\\\ \end{gathered}$

The task is to predict target $\bm{y}\in\mathbb{R}$ based on observed
$\bm{X}\in\mathbb{R}^{20}$, where
$\bm{X}=(\bm{X}_{1}\in\mathbb{R}^{10},\bm{X}_{2}\in\mathbb{R}^{10})$. The
predictor $w\circ\Phi$ is learnt via the IRM objective and as noted
previously, the overparametrization in the objective is handled by
parametrizing $\Phi\in\mathbb{R}^{20}$ to be rank $1$ and fixing $w=1.0$ as a
scalar. The true weights $W_{1},W_{2}\in\mathbb{R}^{10}$ are fixed Gaussian
entries, but the sampling of $c(e)$ for different environments controls
whether $\bm{X}_{2}$ is a causal feature for $\bm{y}$.
The goal is to visualize the intuition in Lemma 1 i.e. how IRM can discard
causal features which are non-invariant. To that end, we sample $1000$ data
points from four environments characterized by $e\in\\{0.2,1,2,5\\}$. The
$c(e)$ for each environment is assigned uniformly, such that the final
training set has two environments for each of $\\{0,1\\}$, and feature
$\bm{X}_{2}$ is non-invariant. The learning procedure for IRM is consistent
with the rest of the paper, with $\lambda=10^{3}$ and the initial $4000$
epochs for annealing the IRM loss.
To study learning of each feature, first denote $\Phi=(\Phi_{1},\Phi_{2})$,
$\Phi_{1},\Phi_{2}\in\mathbb{R}^{10}$. Then note that $\Phi_{i}$ captures the
contribution of feature $\bm{X}_{i}$ in the prediction. We then look at
$\frac{\|\Phi_{i}\|}{\|W_{i}\|}$ (averaged over the random draw of environment
weights). Intuitively, this ratio indicates the information captured by the
IRM predictor for that feature. We visualize the results in Fig. 1, which
demonstrates the tendency for IRM to suppress learning of non-invariant
features.

Figure 1: The plot demonstrates that IRM is incentivized to suppress non-
invariant features, as is the case for feature_2.

###  B.2 Linear Regression

For our choice of learning rate, number of iterations and optimizer and
annealing iterations, we refer to ((Lin, Zhu, and Cui 2022)). While the
reported results were for $\lambda=10^{2}$, we verified similar trends for
$\lambda=10^{3}$.

Hyperparameter | Values
---|---
Number of Iterations | 4000
Learning rate | $10^{-3}$
Optimizer | Adam
IRM Penalty | $10^{2}$
Annealing Iterations | 2000

Table 4: Hyperparameters for experiments on the housing dataset, following
(Lin, Zhu, and Cui 2022)

.

Figure 2: Partitioning vs Conditioning: In the underparametrized linear
regression setting where the number of data points is much greater than
learnable parameters, conditioning is not helpful in terms of improving P-IRM
accuracy and partitioning consistently performs better. Figure 3: For the
regression experiment in the main paper, we find that for both ERM and IRM,
there exists an optimal partition. Note that while ERM consistently finds the
unique optima, the IRM solution has some variance due to the non-convex
objective.

###  B.3 Image Classification

For the image classification experiments on MetaShift (Liang and Zou 2022)
dataset, we follow a similar training pipeline as in (Gulrajani and Lopez-Paz
2020; Liang and Zou 2022). Following (Gulrajani and Lopez-Paz 2020), we
consider ResNet-50 (He et al. 2016), since larger ResNets are known to
generalize better. ResNet-50 was pre-trained on ImageNet (He et al. 2016), and
for domain generalization, the batch normalization and final softmax layers of
ResNet are chopped off (Gulrajani and Lopez-Paz 2020). Then ResNet-50 layers
are followed by non-linear functions, e.g. ReLU, and a final dropout layer
(Gulrajani and Lopez-Paz 2020).

#### Domain Generalization

| Experiment 1 | Experiment 2 | Experiment 3 | Experiment 4 |
---|---|---|---|---|---
| $d=0.17$ | $d=0.54$ | $d=0.81$ | $d=0.92$ | Avg. Performance
P-ERM $(p=25)$ | 0.867(0.050) | 0.740(0.079) | 0.557(0.056) | 0.430(0.079) | 0.6485
P-IRM $(p=25)$ | 0.820(0.148) | 0.742(0.138) | 0.597(0.243) | 0.753(0.209) | 0.728
P-IB_IRM $(p=25)$ | 0.613(0.386) | 0.740(0.171) | 0.203(0.029) | 0.343(0.464) | 0.47475
Table 5: Continued table on Domaing Generalization in Metashift, for $p=25$

Herein, we report the relevant choice of hyperparameters in Table 6, to
reproduce our results pertaining to the domain generalization experiment in
the main paper.

Model |  IRM Penalty weight |  IRM Annealing Iterations |  IB_IRM Penalty weight |  IB_IRM annealing iterations
---|---|---|---|---
Experiment 1 |  |  |  |
IRM | 10 | 20 |  |
P-IRM (p=0) | 10 | 20 |  |
P-IRM (p=10) | 10 | 40 |  |
P-IRM (p=25) | 10 | 20 |  |
IB_IRM | 10 | 40 | 10 | 20
P-IB_IRM (p=0) | 10 | 20 | 10 | 40
P-IB_IRM (p=10) | 100 | 20 | 10 | 20
P-IB_IRM (p=25) | 10 | 40 | 10 | 20
Experiment 2 |  |  |  |
IRM | 10 | 20 |  |
P-IRM (p=0) | 10 | 20 |  |
P-IRM (p=10) | 10 | 20 |  |
P-IRM (p=25) | 10 | 20 |  |
IB_IRM | 10 | 20 | 10 | 20
P-IB_IRM (p=0) | 10 | 20 | 10 | 20
P-IB_IRM (p=10) | 10 | 20 | 10 | 20
P-IB_IRM (p=25) | 10 | 20 | 10 | 20
Experiment 3 |  |  |  |
IRM | 10 | 40 |  |
P-IRM (p=0) | 10 | 40 |  |
P-IRM (p=10) | 10 | 40 |  |
P-IRM (p=25) | 10 | 40 |  |
IB_IRM | 10 | 20 | 10 | 20
P-IB_IRM (p=0) | 10 | 20 | 10 | 40
P-IB_IRM (p=10) | 10 | 20 | 10 | 20
P-IB_IRM (p=25) | 10 | 20 | 10 | 20
Experiment 4 |  |  |  |
IRM | 10 | 20 |  |
P-IRM (p=0) | 10 | 40 |  |
P-IRM (p=10) | 10 | 20 |  |
P-IRM (p=25) | 10 | 20 |  |
IB_IRM | 10 | 20 | 10 | 20
P-IB_IRM (p=0) | 10 | 20 | 10 | 20
P-IB_IRM (p=10) | 10 | 40 | 10 | 20
P-IB_IRM (p=25) | 10 | 40 | 10 | 20
Table 6: Choice of hyper-parameters for Domain Generalization experiments for
IRM and IB_IRM, obtained via train-domain validation.

#### Subpopulation Shift

Following our setup described in main paper, we conduct further experiments to
study performance under subpopulation shifts for the binary classification
task. In subpopulation shifts, communities used for training and testing are
the same, but their relative proportions differ between training and testing
environments, with certain groups often subject to under-representation. The
goal is to obtain a model to do well even for minority groups in the training
data (Koh et al. 2021).
Following (Liang and Zou 2022), the communities are grouped into two
environments: indoor and outdoor. In training, $cat(outdoor)$ and
$dog(indoor)$ subsets are the minority groups, while $cat(indoor)$ and
$dog(outdoor)$ are majority groups. We vary the percentage of minority groups
within the training set to be $m\in\\{0.12,0.01\\}$ of the total training set,
and we keep the size of the training set fixed with 1700 samples. We use a
balanced set of testing by equally sampling from each environment with
balanced labels. Tables 7 and 8 show that the results under subpopulation
shift settings. We report the average accuracy over the four groups, and worst
group accuracy (the group with the worst performace), and the average minority
accuracy which is the average of minority groups in training i.e.
$cat(outdoor)$ and $dog(indoor)$. On the more challenging setting of $m=0.01$
where minority groups are observed minimally, P-IRM models achieve better
worst group and average minority performance. However as expected, it becomes
harder to improve performance for partitioned models as we lower the amount of
available training data, and it is best rely on IRM/ERM.

For reproducibility of results, Table 9 shows the selected hyperparameters.

| $m=0.12$ |  |
---|---|---|---
| Avg. Acc. | Worst Group Acc. | Avg. Minority Acc.
ERM | 0.816(0.209) | 0.722(0.103) | 0.737(0.021)
P-ERM $(p=0)$ | 0.78(0.129) | 0.623(0.157) | 0.675(0.074)
P-ERM $(p=10)$ | 0.76(0.189) | 0.526(0.089) | 0.616(0.127)
P-ERM $(p=25)$ | 0.779(0.154) | 0.590(0.043) | 0.736(0.206)
IRM | 0.638(0.250) | 0.336(0.160) | 0.558(0.314)
P-IRM $(p=0)$ | 0.703(0.158) | 0.560(0.230) | 0.5705(0.015)
P-IRM $(p=10)$ | 0.681(0.190) | 0.518(0.143) | 0.663(0.205)
P-IRM $(p=25)$ | 0.737(0.147) | 0.604(0.227) | 0.7335(0.183)
IB_IRM | 0.639(0.209) | 0.380(0.082) | 0.47(0.127)
P-IB_IRM $(p=0)$ | 0.587(0.136) | 0.451(0.234) | 0.4695(0.026)
P-IB_IRM $(p=10)$ | 0.613(0.301) | 0.264(0.137) | 0.5635(0.424)
P-IB_IRM $(p=25)$ | 0.596(0.178) | 0.426(0.115) | 0.448(0.031)
Table 7: Subpopulation shift on Metashift: The value $m$ represents the portion of minority groups within a training environment. The partitioned models are applied with additional samples up to percentage $p\in\\{0,10,25\\}$ | $m=0.01$ |  |
---|---|---|---
| Avg. Acc. | Worst Group Acc. | Avg. Minority Acc.
ERM | 0.744(0.209) | 0.514(0.113) | 0.559(0.064)
P-ERM $(p=0)$ | 0.747(0.186) | 0.574(0.112) | 0.587(0.018)
P-ERM $(p=10)$ | 0.729(0.254) | 0.481(0.133) | 0.51(0.041)
P-ERM $(p=25)$ | 0.745(0.248) | 0.488(0.013) | 0.532(0.062)
IRM | 0.725(0.185) | 0.509(0.216) | 0.5775(0.097)
P-IRM $(p=0)$ | 0.677(0.232) | 0.426(0.116) | 0.484(0.082)
P-IRM $(p=10)$ | 0.629(0.277) | 0.341(0.260) | 0.6355(0.269)
P-IRM $(p=25)$ | 0.687(0.181) | 0.525(0.189) | 0.531(0.008)
IB_IRM | 0.269(0.179) | 0.470(0.379) | 0.5555(0.121)
P-IB_IRM $(p=0)$ | 0.56(0.27) | 0.208(0.102) | 0.3955(0.265)
P-IB_IRM $(p=10)$ | 0.556(0.206) | 0.366(0.108) | 0.4805(0.162)
P-IB_IRM $(p=25)$ | 0.568(0.138) | 0.398(0.412) | 0.583(0.092)
Table 8: Subpopulation shift: The value $m$ represents the portion of minority groups within a training environment. The partitioned models are applied with additional samples up to percentage $p\in{0,10,25}$ Model |  IRM Penalty weight |  IRM Annealing Iterations |  IB_IRM Penalty weight |  IB_IRM annealing iterations
---|---|---|---|---
Experiment 1 ($m=0.12$ ) |  |  |  |
IRM | 10 | 40 |  |
P-IRM $(p=0)$ | 10 | 40 — 20 |  |
P-IRM $(p=10)$ | 10 | 40 — 20 |  |
P-IRM $(p=25)$ | 10 | 20 — 40 |  |
IB_IRM | 10 | 40 | 10 | 20
P-IB_IRM $(p=0)$ | 10 | 20 | 10 | 20
P-IB_IRM $(p=10)$ | 10 | 20 — 40 | 10 | 40
P-IB_IRM $(p=25)$ | 10 | 20 | 10 | 20 — 40
|  |  |  |
Experiment 2 ($m=0.01$ ) |  |  |  |
IRM | 10 | 40 |  |
P-IRM $(p=0)$ | 1000 — 10 | 40 |  |
P-IRM $(p=10)$ | 100 — 10 | 40 — 20 |  |
P-IRM $(p=25)$ | 100 — 10 | 20 — 40 |  |
IB_IRM | 1000 | 40 | 10 | 40
P-IB_IRM $(p=0)$ | 10 | 40 — 20 | 100 — 10 | 20
P-IB_IRM $(p=10)$ | 100 — 10 | 20 | 100 — 10 | 40 — 20
P-IB_IRM $(p=25)$ | 1000 — 10 | 20 | 100 — 10 | 40 — 20
Table 9: Choice of hyper-parameters for Subpopulation Shifts experiment on
Metashift for IRM and IB_IRM, obtained via train-domain validation. (For
partitioned models since we are using two models, we are reporting the
parameters for both, if they are the same we report one value only)

###  B.4 Language Experiments

For language experiments, NER and TC, we build a classifier based on the pre-
trained language model BERT (Devlin et al. 2019), followed by a dropout and a
linear layer. We also considered DistillBERT (Sanh et al. 2019) and GPT-2
(Radford et al. 2019),but found that BERT-based models outperformed other
networks. We train the models for the maximum number of iterations, (details
in tables 10, and 11) for one seed. Then we select the best number of
iterations to apply for other seeds. The results are the average of three
seeds.

#### Named Entity Recognition (NER)

For the language NER experiments, the best hyperparameter values are reported
in Table 10, and Table 12 which have been selected based on the best model
performance on the validation set. The training was done for 80 epochs, around
which both training and in-domain validation losses stabilize and remain the
same. For annealing epochs, we considered [10, 20, 30, 35, 40] epochs and
found that for all variants of P-IRM/IRM, 40 epochs yielded best performance.
The optimizer and learning rate was based on standard choice for using pre-
trained BERT models.

Hyperparameter | Values
---|---
Maximum Number of epochs | 80
Batch size | 8
Learning rate | $10^{-6}$
Optimizer | Adam
Number of GPUs | 4
Table 10: Hyperparameters choices for experiments on the NER dataset

#### Text Classification (TC)

We consider another language classification task, which identifies the venue
of a published paper 333https://www.semanticscholar.org/product/api, selecting
AAAI and ICML conferences for classification. This task represents a topic
classification task. Our temporal partitions are: 2006-2008, 2009-2011,
2012-2014, and 2015-2017 and we test on papers published between 2018-2020.
The model selection follows as before, training for a maximum of 40 epochs.
The final hyperparameters are reported in Table 14 and 11. Table 13 shows how
partitioning outperforms their baseline methods. P-IRM with two environments
performed the best among all other models.

Hyperparameter | Values
---|---
Maximum Number of epochs | 40
Batch size | 8
Learning rate | $10^{-6}$
Optimizer | Adam
Number of GPUs | 4
Table 11: Hyperparameters choices for the Text Classification task Model |  # envs |  IRM Penalty weight |  IRM annealing iterations |  IB_IRM Penalty weight |  IB_IRM annealing iterations |  #epochs
---|---|---|---|---|---|---
ERM | 4 |  |  |  |  | 44
P-ERM | 3 |  |  |  |  | 54
P-ERM | 4 |  |  |  |  | 58
IRM | 4 | 1000 | 30 |  |  | 53
P-IRM (partitioned) | 3 | 1000 | 40 |  |  | 70
P-IRM (partitioned) | 2 | 1000 | 40 |  |  | 76
P-IRM (conditioned) | 3 | 100 | 40 |  |  | 64
P-IRM (conditioned) | 2 | 100 | 30 |  |  | 66
IB_IRM | 4 | 100 | 40 | 1 | 40 | 57
P-IB_IRM (partitioned) | 3 | 100 | 40 | 1 | 40 | 77
P-IB_IRM (partitioned) | 2 | 100 | 40 | 1 | 40 | 76
P-IB_IRM (conditioned) | 3 | 100 | 40 | 1 | 40 | 76
P-IB_IRM (conditioned) | 2 | 100 | 40 | 1 | 40 | 59
Table 12: NER dataset: Best IRM hyperparameters values selected based on early stopping on validation data Model |  Number of envs |  Training Years |  Testing accuracy (2018-2020)
---|---|---|---
ERM | 4 | 2006-2017 | 0.862 (0.008)
P-ERM | 3 | 2009-2017 | 0.862 (0.004)
P-ERM | 2 | 2012-2017 | 0.875 (0.014)
IRM | 4 | 2006-2017 | 0.846 (0.013)
P-IRM (paritioned) | 3 | 2009-2017 | 0.862 (0.008)
P-IRM (paritioned) | 2 | 2012-2017 | 0.882 (0.016)
P-IRM (conditioned) | 3 | 2009-2017 | 0.869 (0.007)
P-IRM (conditioned) | 2 | 2012-2017 | 0.853 (0.010)
IB_IRM | 4 | 2006-2017 | 0.846 (0.014)
P-IB_IRM (partitioned) | 3 | 2009-2017 | 0.868 (0.001)
P-IB_IRM (partitioned) | 2 | 2012-2017 | 0.874 (0.016)
P-IB_IRM (conditioned) | 3 | 2009-2017 | 0.860 (0.016)
P-IB_IRM (conditioned) | 2 | 2012-2017 | 0.862 (0.011)
Table 13: TC dataset: Results on text classification, comparison between ERM, IRM, IB_IRM and their partitioned variants.  Model |  # envs |  IRM Penalty weight |  IRM Annealing Iterations |  IB_IRM Penalty weight |  IB_IRM annealing iterations |  # epochs
---|---|---|---|---|---|---
ERM | 4 |  |  |  |  | 22
P-ERM | 3 |  |  |  |  | 39
P-ERM | 4 |  |  |  |  | 38
IRM | 4 | 1000 | 20 |  |  | 37
P-IRM (partitioned) | 3 | 1000 | 20 |  |  | 36
P-IRM (partitioned) | 2 | 1000 | 20 |  |  | 37
P-IRM (conditioned) | 3 | 1000 | 20 |  |  | 33
P-IRM (conditioned) | 2 | 1000 | 20 |  |  | 33
IB_IRM | 4 | 1000 | 20 | 0.1 | 20 | 37
P-IB_IRM (partitioned) | 3 | 1000 | 20 | 0.1 | 20 | 36
P-IB_IRM (partitioned) | 2 | 1000 | 20 | 0.1 | 20 | 39
P-IB_IRM (conditioned) | 3 | 1000 | 20 | 0.1 | 20 | 33
P-IB_IRM (conditioned) | 2 | 1000 | 20 | 0.1 | 20 | 33
Table 14: TC dataset: Best IRM hyperparameters values selected based on early
stopping on validation data