A SAMPLE COMPLEXITY PERSPECTIVE
Year: 2021

Abstract:
Recently, invariant risk minimization (IRM) was proposed as a promising solution to address out-of-distribution (OOD) generalization. However, it is unclear when IRM should be preferred over the widely-employed empirical risk minimization (ERM) framework. In this work, we analyze both these frameworks from the perspective of sample complexity, thus taking a firm step towards answering this important question. We find that depending on the type of data generation mechanism, the two approaches might have very different finite sample and asymptotic behavior. For example, in the covariate shift setting we see that the two approaches not only arrive at the same asymptotic solution, but also have similar finite sample behavior with no clear winner. For other distribution shifts such as those involving confounders or anti-causal variables, however, the two approaches arrive at different asymptotic solutions where IRM is guaranteed to be close to the desired OOD solutions in the finite sample regime for polynomial generative models, while ERM is biased even asymptotically. We further investigate how different factors — the number of environments, complexity of the model, and IRM penalty weight — impact the sample complexity of IRM in relation to its distance from the OOD solutions.

heading: 1 INTRODUCTION
text: A recent study shows that models trained to detect COVID-19 from chest radiographs rely on spurious factors such as the source of the data rather than the lung pathology (DeGrave et al., 2020). This is just one of many alarming examples of spurious correlations failing to hold outside a specific training distribution. In one commonly cited example, Beery et al. (2018) trained a convolutional neural network (CNN) to classify camels from cows. In the training data, most pictures of the cows had green pastures, while most pictures of camels were in the desert. The CNN picked up the spurious correlation and associated green pastures with cows thus failing to classify cows on beaches.
Recently, Arjovsky et al. (2019) proposed a framework called invariant risk minimization (IRM) to address the problem of models inheriting spurious correlations. They showed that when data is gathered from multiple environments, one can learn to exploit invariant causal relationships, rather than relying on varying spurious relationships, thus learning robust predictors. More recent work suggests that empirical risk minimization (ERM) is still state-of-the-art on many problems requiring OOD generalization (Gulrajani & Lopez-Paz, 2020). This gives rise to a fundamental question: when is IRM better than ERM (and vice versa)? In this work, we seek to answer this question through a systematic comparison of the sample complexity of the two approaches under different types of train and test distributional mismatches.
The distribution shifts ( Ptrain(X,Y ) 6= Ptest(X,Y ) ) that we consider informally stated satisfy an invariance condition – there exists a representation Φ∗ of the covariates such that Ptrain(Y |Φ∗(X)) = Ptest(Y |Φ∗(X)) = P(Y |Φ∗(X)). A special case of this occurs when Φ∗ is identity – Ptrain(X) 6= Ptest(X) but Ptrain(Y |X) = Ptest(Y |X) – such a shift is known as a covariate-shift (Gretton et al., 2009). In many other settings Φ∗ may not be identity (denoted as I), examples include settings with confounders or anti-causal variables (Pearl, 2009) where covariates appear spuriously correlated with the label and Ptrain(Y |X) 6= Ptest(Y |X). We use causal Bayesian networks to illustrate these shifts in Figure 1. Suppose Xe = [Xe1 , X e 2 ] represents the image, where X e 1 is the shape of the animal and Xe2 is the background color, Y e is the label of the animal, and e is the index of the environment/domain. In Figure 1a) Xe2 is independent of (Y e, Xe1), it represents the covariate shift
case (Φ∗ = I). In Figure 1b) Xe2 is spuriously correlated with Y e through the confounder εe. In Figure 1c) Xe2 is spuriously correlated with Y e as it is anti-causally related to Y e. In both Figure 1b) and c) Φ∗ 6= I; Φ∗ is a block diagonal matrix that selects Xe1 . Our setup assumes we are given data from multiple training environments satisfying the invariance condition, i.e., P(Y |Φ∗(X)) is the same across all of them. Ideally, we want to learn and predict using E[Y |Φ∗(X)]; this predictor has a desirable OOD behavior as we show later where we prove min-max optimality with respect to (w.r.t.) unseen test distributions satisfying the invariance condition. Our goal is to analyze and compare ERM and IRM’s ability to learn E[Y |Φ∗(X)] from finite training data acquired from a fixed number of training environments. Our analysis has two parts.
1) Covariate shift case (Φ∗ = I): ERM and IRM achieve the same asymptotic solution E[Y |X]. We prove (Proposition 4) that the sample complexity for both the methods is similar thus there is no clear winner between the two in the finite sample regime. For the setup in Figure 1a), both ERM and IRM learn a model that only uses Xe1 .
2) Confounder/Anti-causal variable case (Φ∗ 6= I): We consider a family of structural equation models (linear and polynomial) that may contain confounders and/or anti-causal variables. For the class of models we consider, the asymptotic solution of ERM is biased and not equal to the desired E[Y |Φ∗(X)]. We prove that IRM can learn a solution that is within O( √ ) distance from E[Y |Φ∗(X)] with a sample complexity that increases as O( 1 2 ) and increases polynomially in the complexity of the model class (Proposition 5, 6); (defined later) is the slack in IRM constraints. For the setup in Figure 1b) and c), IRM gets close to only using Xe1 , while ERM even with infinite data (Proposition 17 in the supplement) continues to use Xe2 . We summarize the results in Table 1.
Arjovsky et al. (2019) proposed the colored MNIST (CMNIST) dataset; comparisons on it showed how ERM-based models exploit spurious factors (background color). The CMNIST dataset relied on anti-causal variables. Many supervised learning datasets may not contain anti-causal variables (e.g. human labeled images). Therefore, we propose and analyze three new variants of CMNIST in addition to the original one that map to different real-world settings: i) covariate shift based CMNIST (CS-CMNIST): relies on selection bias to induce spurious correlations, ii) confounded CMNIST (CF-CMNIST): relies on confounders to induce spurious correlations, iii) anti-causal CMNIST (ACCMNIST): this is the original CMNIST proposed by Arjovsky et al. (2019), and iv) anti-causal and confounded (hybrid) CMNIST (HB-CMNIST): relies on confounders and anti-causal variables to induce spurious correlations. On the latter three datasets, which belong to the Φ∗ 6= I class described above, IRM has a much better OOD behavior than ERM, which performs poorly regardless of the data size. However, IRM and ERM have a similar performance on CS-CMNIST with no clear winner. These results are consistent with our theory and are also validated in regression experiments.

heading: 2 RELATED WORKS
text: IRM based works. Following the original work IRM from Arjovsky et al. (2019), there have been several interesting works — (Teney et al., 2020; Krueger et al., 2020; Ahuja et al., 2020; Chang et al., 2020; Mahajan et al., 2020) is an incomplete representative list — that build new methods inpired from IRM to address the OOD generalization problem. Arjovsky et al. (2019) prove OOD guarantees for linear models with access to infinite data from finite environments. We generalize these results in several ways. We provide a first finite sample analysis of IRM. We characterize the impact of hypothesis class complexity, number of environments, weight of IRM penalty on the sample complexity and its distance from the OOD solution for linear and polynomial models.
Theory of domain generalization and domain adaption. Following the seminal works (BenDavid et al., 2007; 2010), there have been many interesting works — (Muandet et al., 2013; Ajakan et al., 2014; Zhao et al., 2019; Albuquerque et al., 2019; Li et al., 2017; Piratla et al., 2020; Matsuura & Harada, 2020; Deng et al., 2020; David et al., 2010; Pagnoni et al., 2018) is an incomplete representative list (see Redko et al. (2019) for further references) — that build the theory of domain adaptation and generalization and construct new methods based on it. While many of these works develop bounds on loss over the target domain using train data and unlabeled target data,
some (Ben-David & Urner, 2012; David et al., 2010; Pagnoni et al., 2018) analyze the finite sample (PAC) guarantees for domain adaptation under covariate shifts. These works (Ben-David & Urner, 2012; David et al., 2010; Pagnoni et al., 2018) access unlabeled data from a target domain, which we do not. Instead, we have data from multiple training domains (as in domain generalization). In these works, the guarantees are w.r.t. a specific target domain, while we provide (for linear and polynomial models) worst-case guarantees w.r.t. all the unseen domains satisfying the invariance condition. Also, we consider a larger family of distribution shifts including covariate shifts. The above two categories are not exhaustive – e.g., there are some recent works that characterize how some inductive biases favor extrapolation Xu et al. (2021) and can be better for OOD generalization.

heading: 3 SAMPLE COMPLEXITY OF INVARIANT RISK MINIMIZATION

heading: 3.1 INVARIANT RISK MINIMIZATION
text: We start with some background on IRM (Arjovsky et al., 2019). Consider a datasetD = {De}e∈Etr , which is a collection of datasets De = {(xei , yei , e)} ne i=1} obtained from a set of training environments Etr, where e is the index of the environment, i is the index of the data point in the environment, ne is the number of points from environment, xei ∈ X ⊆ Rn is the feature value and yei ∈ Y ⊆ R is the corresponding label. Define a probability distribution {πe}e∈Etr , πe is the probability that a training data point is from environment e. Define a probability distribution of points conditional on environment e as Pe, (Xe, Y e) ∼ Pe. Define the joint distribution P̄, (Xe, Y e, e) ∼ P̄, dP̄(Xe, Y e, e) = πedPe(Xe, Y e). D is a collection of i.i.d. samples from P̄. Define a predictor f : X → R and the space F of all the possible maps from X → R. Define the risk achieved by f in environment e as Re(f) = Ee [ ` ( f(Xe), Y e )] , where ` is the loss, f(Xe) is the predicted value, Y e is the corresponding label and Ee is the expectation conditional on environment e. The overall expected risk across the training environments is R(f) = ∑ e∈Etr π
eRe(f). We are interested in two settings: regression (square loss) and binary-classification (cross-entropy loss). In the main body, our focus is regression (square loss) and we mention wherever the results extend to binary-classification (cross-entropy). We discuss these extensions in the supplement.
OOD generalization problem. We want to construct a predictor f that performs well across many unseen environments Eall, where Eall ⊇ Etr. For o ∈ Eall\Etr, the distribution Po can be very different from the train environments. Next we state the OOD problem.
min f∈F max e∈Eall
Re(f) (1)
The above problem is very challenging to solve since we only have access to data from training environments Etr but are required to find the robust solution over all environments Eall. Next, we make assumptions on Eall and characterize the optimal solution to equation 1. Assumption 1. Invariance condition. There exists a representation Φ∗ that transforms Xe to Ze = Φ∗(Xe) and ∀e, o ∈ Eall,∀z ∈ Φ∗(X ) satisfies Ee[Y e|Ze = z] = Eo[Y o|Zo = z]. Also, ∀e ∈ Eall,∀z ∈ Φ∗(X ), Vare[Y e|Ze = z] = ξ2, where Vare is the conditional variance.
The above assumption is inspired from causality (Pearl, 2009). Φ∗ acts as the causal feature extractor and from the definition of causal features, it follows that Ee[Y e|Ze = z] does not vary across
environments. When a human labels a cow she uses Φ∗ to extract causal features from the pixels to identify cow while ignoring the background. The first part of the above assumption encompasses a large class of distribution shifts including standard covariate shifts (Gretton et al., 2009). Covariate shift assumes ∀e, o ∈ Eall,∀x ∈ X , P(Y e|Xe = x) and P(Y o|Xo = x) are equal thus implying Ee[Y e|Xe = x] = Eo[Y o|Xo = x]. Therefore, for covariate shifts, Φ∗ is identity in Assumption 1. A simple instance illustrating Assumption 1 with Φ∗ = I is when Y e ← g(Xe)+εe, where Ee[εe] = 0,Ee[(εe)2] = σ2, εe ⊥ Xe. Using Assumption 1, we define the invariant map m : Φ∗(X )→ R as follows ∀z ∈ Φ∗(X ), m(z) = Ee[Y e|Ze = z], where Ze = Φ∗(Xe) (2) Assumption 2. Existence of an environment where the invariant representation is sufficient. ∃ an environment e ∈ Eall such that Y e ⊥ Xe|Ze
Assumption 2 states there exists an environment where the information that Xe has about Y e is also contained in Ze. Define a composition m ◦ Φ∗, ∀x ∈ X , m ◦ Φ∗(x) = Ee[Y e|Ze = Φ∗(x)]. Proposition 1. If ` is the square loss, and Assumptions 1 and 2 hold, then m ◦ Φ∗ solves the OOD problem (equation 1).
The proofs of all the propositions are in the supplement. A similar result holds for the cross-entropy loss (discussion in supplement). For the rest of the paper, we focus on learning m ◦ Φ∗ as it solves the OOD problem. For covariate shifts Φ∗ = I, m(x) = Ee[Y e|Xe = x] is the OOD solution. In Arjovsky et al. (2019), a proof connecting m ◦ Φ∗ and OOD was not stated. Recently, in Koyama & Yamaguchi (2020), a result similar to Proposition 1 was shown but with a few differences. The authors assume conditional probabilities are invariant unlike our assumption that only requires conditional expectations and variances to be invariant. However, their result applies to more losses. m ◦Φ∗ is the target we want to learn. Arjovsky et al. (2019) proposed IRM since standard min-max optimization over the training environments Etr and ERM fail to learn m ◦ Φ∗ in many cases. The authors in Arjovsky et al. (2019) identify a crucial property of m ◦ Φ∗ and use it to define an object called invariant predictor that we define next.
Invariant predictor and IRM optimization. Define a representation map Φ : X → Z from feature space to representation space Z ⊆ Rq . Define a classifier map, w : Z → R from representation space to real values. DefineHΦ andHw as the spaces of representations and classifiers respectively. A data representation Φ elicits an invariant predictor w ◦ Φ across environments e ∈ Etr if there is a classifier w that achieves the minimum risk simultaneously for all the environments, i.e., ∀e ∈ Etr, w ∈ arg minw̄∈Hw Re(w̄ ◦ Φ). Observe that if we we transform the data with representation Φ∗ then m will achieve the minimum risk simultaneously in all the environments. If Φ∗ ∈ HΦ and m ∈ Hw, then m ◦ Φ∗ is an invariant predictor. IRM selects the invariant predictor with least sum risk across environments (results presented later can be adapted if invariant predictor was selected based on the worst-case risk over the environments as well) as follows:
min Φ∈HΦ,w∈Hw R(w ◦ Φ) = ∑ e∈Etr πeRe(w ◦ Φ)
s.t. w ∈ arg min w̄∈Hw
Re(w̄ ◦ Φ), ∀e ∈ Etr (3)
From the above discussion we know m ◦ Φ∗ is a feasible solution to equation 3. It is also the ideal solution we want IRM to find since it solves equation 1. Later in Propositions 4, 5, and 6, we show that IRM actually solves equation equation 1. For the setups in Proposition 5, and 6, conventional ERM based approaches fail thus justifying the need for above formulation.

heading: 3.2 SAMPLE COMPLEXITY OF GRADIENT CONSTRAINT FORMULATION OF IRM
text: In Arjovsky et al. (2019), a gradient constrained alternate (derived below in equation 4) to equation 3 was proposed, which focuses on linear and scalar classifiers (Z = R, Φ : X → R, Hw = R). In this case, the composite predictor w ◦Φ is a multiplication of w and Φ written as w ·Φ. (For binaryclassification predictor’s outputw·Φ(x) represents logits.) From the definition of invariant predictors and Hw = R it follows that if ∀w̄ ∈ R, Re(1 · Φ) ≤ Re(w̄ · Φ), then Φ is an invariant predictor. For square and cross-entropy losses, Re(w · Φ) is convex in w. Therefore, a gradient constraint ‖∇w|w=1.0Re(w·Φ)‖ = 0 is equivalent to the condition that ∀w̄ ∈ R, Re(1·Φ) ≤ Re(w̄ ·Φ), which implies Φ is an invariant predictor. Recall that IRM aims is to search among invariant predictors and find one that minimizes the risk. We state this as a gradient constrained optimization as follows
min Φ∈HΦ R(Φ) s.t. ∥∥∇w|w=1.0Re(w · Φ)∥∥ = 0, ∀e ∈ Etr (4)
We propose an approximation of the above with slack in the constraint. Define R ′ (Φ) =∑
e∈Etr π e‖∇w|w=1.0Re(w · Φ)‖2 and a set S IV( ) = {Φ | R
′ (Φ) ≤ ,Φ ∈ HΦ}. Note that R ′ is
very similar to the penalty defined in Arjovsky et al. (2019). The approximation of equation 4 is
min Φ∈S IV( ) R(Φ) (5)
If = 0, then equation 4 and equation 5 are equivalent. In all the optimizations so far, the expectations are computed w.r.t. the distributions Pe, which are unknown. Therefore, we develop an empirical version of equation 5 below (in equation 6) and call it empirical IRM (EIRM). We replace R and R ′ with empirical estimators R̂ and R̂ ′ respectively. For R we use a simple plugin estimator (sample mean of loss across all the samples in D). For R ′
we construct a new estimator that enables the use of standard concentration inequalities. Define a set Ŝ IV( ) = {Φ | R̂′(Φ) ≤ ,Φ ∈ HΦ}.
min Φ∈Ŝ IV( ) R̂(Φ) (6)
If we replaced Ŝ IV( ) with HΦ in equation 6, then we get the standard ERM. ERM aims to solve minΦ∈HΦ R̂(Φ). The sample complexity analysis of ERM aims to understand the distance between the empirical solutions and the expected solutions as a function of the number of samples. Similarly, we seek to understand the relationship between solutions of equation 6 and equation 5. Assumption 3. Bounded loss and bounded gradient of the loss. ∃ L < ∞, L′ < ∞ such that ∀Φ ∈ HΦ,∀x ∈ X ,∀y ∈ Y, |`(Φ(x), y)| ≤ L, |∂`(w·Φ(x),y)∂w |w=1.0| ≤ L ′ .
If every Φ in the hypothesis class HΦ is bounded by M and the label space Y is bounded, then for both square and cross-entropy loss, `(Φ(·), ·) and ∂`(w·Φ(·),·)∂w |w=1.0 are bounded. Define κ = minΦ∈HΦ |R ′ (Φ) − |; κ measures how close any penalty can get to the boundary . κ quantifies how good the finite sample approximation R̂ ′
need to be in order to get Ŝ IV( ) = S IV( ). Define ν to quantify the approximation w.r.t. optimal risk. Proposition 2. For every ν > 0, > 0 and δ ∈ (0, 1), ifHΦ is a finite hypothesis class, Assumption 3 holds, κ > 0, and if the number of samples |D| is greater than max { 16L′4
κ2 , 8L2 ν2
} log ( 4|HΦ|
δ
) ,
then with a probability at least 1 − δ, every solution Φ̂ to EIRM (equation 6) is a ν approximation of IRM, i.e. Φ̂ ∈ S IV( ), R(Φ∗) ≤ R(Φ̂) ≤ R(Φ∗) + ν, where Φ∗ is a solution to IRM (equation 5).
Proof Sketch. The standard analysis in learning theory on ERM or regularized/constrained ERM typically relies on linearly separable loss functions. In such cases, we can use standard plugin estimators and analyze their behavior using concentration inequalities. In our setting, R ′ is a weighted sum of squares of expectation and thus it is not linearly separable. We develop a new way of expressing R ′ that allows us to make it linearly separable. Next, in order to ensure R(Φ∗) ≤ R(Φ̂) ≤ R(Φ∗) + ν we first need to guarantee that the set of invariant predictors is exactly recovered, i.e., Ŝ IV( ) = S IV( ) (exact recovery is typically not required in existing constrained analysis such as Woodworth et al. (2017) Agarwal et al. (2018)). We show that if the number of samples grow as 1κ2 even the closest points on either side of the boundary of the set S
IV( ) are correctly discriminated, which guarantees exact recovery of S IV( ). Once the exact set is recovered, beyond this we use standard learning theory tools to ensure R(Φ∗) ≤ R(Φ̂) ≤ R(Φ∗) + ν. The above result holds for both square and cross-entropy loss. For ease of exposition, we use the standard setting of finite hypothesis class and extend all the results to infinite hypothesis classes in the supplement (summary of insights from the extension are in Section 3.3.2). Next, we state a standard result on ERM’s sample complexity. Define a Φ+ such that Φ+ ∈ arg minΦ∈HΦ R(Φ) Proposition 3. (Shalev-Shwartz & Ben-David, 2014) For every ν > 0 and δ ∈ (0, 1), if HΦ is a finite hypothesis class, Assumption 3 holds, and if the number of samples |D| is greater than 8L2
ν2 log ( 2|HΦ| δ ) , then with a probability at least 1 − δ, every solution Φ† to ERM is an ν approxi-
mation of expected risk minimization, i.e., R(Φ+) ≤ R(Φ†) ≤ R(Φ+) + ν.
Proposition 2 vs. 3 Since κ ≤ , the sample complexity of EIRM grows at least asO(max{ 1 2 , 1 ν2 }). Let us look at the two terms inside max- i) 1ν2 growth term is similar to ERM, it ensures ν approximate optimality in the overall risk R, ii) 1 2 growth ensures the IRM penalty R ′ is less than . A direct comparison of sample complexities in Propositions 2 and 3 suggests that the sample complexity of EIRM is higher than ERM, which is not the complete picture. The two approaches may not converge to the same solutions and IRM may converge to a solution with better OOD behavior than one achieved by ERM. Therefore, a fair comparison is only possible when we also study the OOD properties of the solutions achieved by the two approaches, which is the subject of the next section.

heading: 3.3 OOD PERFORMANCE: ERM VS. IRM
text: We divide the comparisons based on distributional shift assumptions that decide whether ERM and IRM arrive at the same asymptotic solutions or not.

heading: 3.3.1 COVARIATE SHIFT
text: Assumption 4. Invariance w.r.t. all the features. ∀e, o ∈ Eall and ∀x ∈ X , E[Y e|Xe = x] = E[Y o|Xo = x]. ∀e ∈ Eall, Xe ∼ PeXe and the support of PeXe is equal to X .
As stated earlier, the first part of the above assumption follows from standard covariate shift assumptions (Gretton et al., 2009) and is a special case of the first part of the Assumption 1 with Φ∗ set to I. If Φ∗ = I, then m (equation 2), which simplifies to m(x) = E[Y e|Xe = x], solves the OOD problem equation 1. A generative model that satisfies the above Assumption 4 is given as
Y e ← g(Xe) + εe, E[εe] = 0, εe ⊥ Xe, E[(εe)2] = σ2 (7) In the above model Xe is the cause, Y e is the effect, and g a general non-linear function (it satisfies Assumption 4 with m = g). Next, we compare ERM and IRM’s ability to learn m under covariate shifts. In Figure 1 a), we show Define κ̃ = minΦ1,Φ2∈HΦ,Φ1 6=Φ2 |R(Φ1)−R(Φ2)|, which measures the minimum separation between the risks of any two distinct hypothesis inHΦ. Proposition 4. Let ` be the square loss. For every ν > 0, > 0 and δ ∈ (0, 1), if HΦ is a finite hypothesis class, m ∈ HΦ, Assumptions 3, 4 hold, and
• if the number of samples |D| is greater than max { 8L2
ν2 log( 4|HΦ| δ ), 16L′4 2 log( 2 δ ) }
, then with a probability at least 1−δ, every solution Φ̂ to EIRM (equation 6) satisfiesR(m) ≤ R(Φ̂) ≤ R(m)+ ν. If also ν < κ̃, then Φ̂ = m.
• if the number of samples |D| is greater than 8L 2 ν2 log( 2|HΦ| δ ), then with a probability at least 1− δ, every solution Φ† to ERM satisfies R(m) ≤ R(Φ†) ≤ R(m) + ν. If also ν < κ̃, then Φ† = m.
Implications of Proposition 4. ERM and EIRM both asymptotically achieve the ideal OOD solution; the above proposition helps compare them in a finite sample regime. The second term inside the max for EIRM, 16L ′4
2 log( 2 δ ), does not depend on the size of the hypothesis class. Hence, for
large hypothesis classes, the sample complexity of EIRM equals 8L 2 ν2 log( 4|HΦ| δ ). Consequently, the sample complexity of EIRM and ERM differs by a constant 8L 2
ν2 log(2). Thus we conclude, for large hypothesis classes, both ERM and EIRM have similar sample complexity. Next, we contrast the sample complexity of EIRM in Proposition 4, O( 1ν2 ), to Proposition 2, O(max{ 1 2 , 1 ν2 }); the additional covariate shift assumption in Proposition 4 helps get to a lower sample complexity of O( 1ν2 ). In Proposition 4, we assumed square loss, but a similar result extends to cross-entropy loss as well.

heading: 3.3.2 DISTRIBUTIONAL SHIFT WITH CONFOUNDERS AND (OR) ANTI-CAUSAL VARIABLES
text: In this section, we consider more general models than equation 7, which only contained cause Xe and effect Y e. We also allow confounders and anti-causal variables. However, we restrict g to polynomials. We start with linear models from Arjovsky et al. (2019). Assumption 5.
e ∼ Categorical({πo}o∈Etr ), ∀o ∈ Etr, πo > 0 Y e ← γT(Ze1) + εe, εe ⊥ Ze1 , E[εe] = 0, E[(εe)2] = σ2, |εe| ≤ εsup
Xe ← S(Ze1 , Ze2) (8)
We assume that Z1 component of S is invertible, i.e. ∃ S̃ such that S̃(S(Z1, Z2)) = Z1, and γ 6= 0. ∀e ∈ Etr, πe ≥ π min
|Etr| > 0. Define Σ e = E[XeXe,T]. ∀e ∈ Etr, Σe is positive definite. The support
of distribution of Ze = (Ze1 , Z e 2), PeZe , is bounded and the norm of S, ‖S‖ = σmax(S) (maximum singular value of S), is also bounded.
In the above model, Ze1 is the cause of X e and Y e but may not be directly observed. Ze2 may be arbitrarily correlated with Ze1 and e. We observe a scrambled transformation of (Ze1 , Z e 2) in Xe. If Ze2 is an effect of Y e (Ze2 ← Y e + Ne), then Ze2 is an anti-causal variable. If He causes both εe (εe ← He + N̄e) and Ze2 (Ze2 ← He + Ñe), then He is a confounder. In both these cases, Ze2 is spuriously correlated with the label Y
e. Consequently, the standard ERM based models estimate Ze2 from X e and end up being biased w.r.t the desired OOD model, which does not use
Ze2 . If Assumptions 5, 2 hold, then a linear model S̃ Tγ (Xe predict−−−−→ γTS̃Xe = γTZe1 ) solves the OOD problem in equation 1 (Φ∗ = S̃, m = γT in Proposition 1) and it relies only on Ze1 . In the supplement (Proposition 17), we prove that ERM based models do not recover S̃Tγ. Assumption 6. Linear general position. A set of training environments Etr is said to lie in a linear general position of degree r for some r ∈ N if |Etr| > n− r + n/r and for all non-zero x ∈ Rn
dim ( span { Ee[XeXe,T]x− Ee[Xeεe] } e∈Etr ) > n− r (9)
where span is the linear span, dim is the dimension, and recall n is dimension of Xe. This assumption checks for diversity in the environments and holds almost everywhere (Arjovsky et al., 2019). Assumption 7. Inductive bias. HΦ is a finite set of linear models parametrized by Φ ∈ Rn (output ΦTXe). S̃Tγ ∈ HΦ. ∃ ω > 0,Ω > 0, s.t. ∀Φ ∈ HΦ, ω ≤ ‖Φ‖2 ≤ Ω & 2ω ≤ ‖S̃Tγ‖2 ≤ 23+2√2Ω.
Informally stated, the above assumption requires the OOD optimal predictor S̃Tγ to lie in the interior of the search space and not on the boundary. If Assumptions 5, 7 hold, then Assumption 3 holds. Hence, we can use the bounds L and L′ on ` and ∂`(w·Φ(·),·)∂w |w=1.0 respectively in our next result. Define the minimum eigenvalue across all Σe as λmin = mine∈Etr λmin(Σe),
th = (24−16
√ 2)
3 πmin |Etr| (ωλmin) 2 and τ = 12ωλmin √ 3|Etr| 2πmin
. Next, we analyze how EIRM learns S̃Tγ. Proposition 5. Let ` be the square loss. For every ∈ (0, th) and δ ∈ (0, 1), if Assumptions 5, 6 (with r = 1), 7 hold and if the number of data points |D| is greater than 16L ′4 2 log ( 2|HΦ|
δ
) , then with
a probability at least 1 − δ, every solution Φ̂ to EIRM (equation 6) satisfies Φ̂ = (S̃Tγ ) α, where
α ∈ [ 1 1+τ √ , 1 1−τ √ ].
Proof Sketch. In learning theory it is common to analyze the concentration of empirical risks around the expected risks. In our case, we have a target ideal solution to equation 1 (S̃Tγ) and we want our empirical solutions to concentrate around that. A direct finite sample approximation of equation 4 is hard to analyze. Therefore, we introduce an intermediate problem in equation 5 and then develop a finite sample approximation of it in equation 6. We first show that solving equation 5 leads to solutions in the neighborhood of the target. To show this we use the linear general position assumption. Next, we connect equation 6 and equation 5 using our new estimator for R ′ and Hoeffding’s inequality.
Implications of Proposition 5. 1. Convergence rate of ERM vs. EIRM: Recall that is the slack on IRM penalty R ′ . If is sufficiently small and the data grows asO( 1 2 ), every solution Φ̂ to EIRM (equation 6) is in √ radius of the OOD solution, i.e., ‖Φ̂− S̃Tγ‖ = O( √ ). We contrast these rates to ones in the covariate shift setting (Section 3.3.1). Let E[Y e|Xe = x] = ΨTx. If the data grows as O( 1ν2 ), then both ERM and EIRM solution converge to Ψ as ‖Φ̂−Ψ‖ = O( √ ν) (from Proposition 4). This shows that EIRM works in more settings (Proposition 4, 5) than ERM while matching the convergence rate of ERM.
2. Sample complexity grows polynomially in data dimension to ensure OOD generalization: Next, we set = µ th with µ ∈ [0, 1), and |Etr| = 2n (satisfies Assumption 6 for r = 1). A simple manipulation of terms in Proposition 5 shows that sample complexity with quadratic growth in data dimension n, O ( n2 µ2 log ( 2|HΦ| δ )) , ensures Φ̂ = (S̃Tγ ) α with α ∈ [ 1 1+ √ µ( √ 2−1) , 1 1−√µ( √ 2−1) ].
3. Comparison with Proposition 2: Lastly, we contrast sample complexity of EIRM in Proposition 5, O( 1 2 ), to Proposition 2, O(max{ 1 2 , 1 ν2 }); the additional distributional assumptions in Proposition 5 help arrive at a lower sample complexity of O( 1 2 ). The bound in Proposition 2, O(max{ 1 2 , 1 ν2 }), is larger than the one in Proposition 4, O( 1 ν2 ), and Proposition 5, O( 1 2 ), but is more general as it is agnostic to the distributional assumptions.
A simple illustration summarizing Propositions 4, 5: Set S to identity in Assumption 5. Recall Ze1 and Z e 2 from Assumption 5. Since S is identity X e can be written as [Xe1 , X e 2 ], where X e 1 = Z e 1 and Xe2 = Z e 2 . If X e 2 ⊥ εe, then E[Y e|Xe] is invariant and Assumption 4 holds. This corresponds to the setup in Figure 1a). We can now use Proposition 4 and deduce that ERM and IRM have same sample complexity and end up learning the ideal model that only uses the causal features Xe1 . If Xe2 ← εe +Ne, then this corresponds to the setup Figure 1b), Xe1 is the cause and Xe2 is spuriously correlated with label Y e through the confounder εe. If Xe2 ← Y e+Ne, then this corresponds to the setup in Figure 1c), Xe1 is the cause and X e 2 is anti-causally related to the label Y
e. In both these cases, the ideal OOD solution that solves equation 1 will only exploitXe1 to make predictions. From Propositions 5, it follows that IRM when fed with O( 1 2 ) samples, it is in √ radius of the target OOD solution, while ERM is asymptotically biased and exploits Xe2 (Proposition 17). We define a polynomial version of the model in Assumption 5. We only need to change Y e ← γT(Ze1) + ε
e to Y e ← γtζp(Ze1) + εe. ζp is a polynomial feature map of degree p defined as ζp : Rc → Rc ′ , where c is the dimension of the input Ze1 , ζp(W ) = ( W,W ⊗ W, . . . , (W ⊗
W...p times ⊗ W ) ) = ( (W⊗i)pi=1 ) and ⊗ is the Kronecker product. Also, c′ = ∑p i=1 c
i. Can we directly use the analysis from the linear case by transforming Xe appropriately? No, we first need to find an appropriate transformation for the scrambling matrix S that satisfies the conditions (invertibiltiy) in Assumption 5 while maintaining a linear relationship between transformations of Xe and Ze. We present the main result informally below (details are in the supplement). Proposition 6. (Informal statement) For sufficiently small and δ ∈ (0, 1), if Assumptions similar to Proposition 5 hold and |D| ≥ 16L ′4
2 log( 2|HΦ| δ ), then with a probability at least 1−δ, every solution
Φ̂ to EIRM (equation 6) satisfies Φ̂ = ¯̃STγ(α), where ¯̃STγ is the OOD optimal model (defined in the supplement), α ∈ [ 1
1+τ √ , 1 1−τ √ ].
Insights from the polynomial case and infinite hypothesis case. In the polynomial case, we adapt the linear general position Assumption 6, the number of environments |Etr| are now required to grow as O(np). As a result, in the sample complexity analysis we discussed, we replace n with np to obtain that a sample complexity of O ( n2p µ2 log ( 2|HΦ| δ )) ensures Φ̂ = ¯̃STγ(α) with α ∈ [ 1 1+ √ µ( √ 2−1) , 1 1−√µ( √ 2−1) ]. In the infinite hypothesis case, the main change in the results is that we replace |HΦ| with an appropriate model complexity metric (Shalev-Shwartz & Ben-David, 2014). Consider Proposition 5, a sample complexity of O ( n3 µ2 log ( n µ )) ensures Φ̂ = (S̃Tγ ) α with
α ∈ [ 1 1+ √ µ( √ 2−1) , 1 1−√µ( √ 2−1) ] in contrast toO ( n2 µ2 log ( 2|HΦ| δ )) in the finite hypothesis case. We showed the benefits of IRM for polynomial models and other extensions (non-linear S) are future work. In the supplement, we provide a dialogue explaining how our work fits in the big picture.

heading: 4 EXPERIMENTS
text: In this section, we discuss classification experiments (regression experiments with similar qualitative findings are in the supplement). We introduce three new variants of the colored MNIST (CMNIST) dataset in (Arjovsky et al., 2019). We divide the training data in MNIST digits into two environments (e = 1, 2) equally and the testing data in MNIST digits is assigned to another environment (e = 3). Xeg : gray scale image of the digit, Y e g : label of the gray scale digit (digits≥ 5 have Y eg = 1 and digits < 5 have Y eg = 0). X e: final colored image and Y e: final label are generated as follows. Define Bernoulli variables G, N , Ne that take a value 1 with probability θ, β and βe and 0 otherwise. Define a color variable Ce, where Ce = 0 is red and Ce = 1 is green. Let ⊕ denotes xor operation. Y eg ← L(Xeg ), L : Human labeling, Y e ← L(Xeg )⊕N, Corrupt the original labels with noise Ce ← G(Y e ⊕Ne) + (1−G)(N ⊕Ne), Use G to select b/w anti-causal or confounded Xe ← T (Xg, Ce), T : transformation to color the image (10)
If the probability θ = 1, then G = 1 and that gives us back the original CMNIST in (Arjovsky et al., 2019), which we call anti-causal CMNIST (AC-CMNIST). If θ = 0, then we get confounded colored MNIST (CF-CMNIST). If 0 < θ < 1, we get a hybrid dataset (HB-CMNIST). The above model in equation 10 has features of the model in Assumption 5, where Xeg , C
e, L, T take the role of Ze1 , Z e 2 , γ, S. We set the noise N parameter β = 0.25, and the parameter for N
e in the three environments [β1, β2, β3] = [0.1, 0.2, 0.9]. Color is spuriously correlated with the label; P(Ce = 1|Y e = 1) varies drastically in the three environments ([0.9, 0.8, 0.1] for AC-CMNIST). In the variants of CMNIST we discussed, P(Y e|Xe) varies across the environments. We now define a covariate shift based CMNIST (CS-CMNIST) (P(Y e|Xe) is invariant). We use selection bias to induce spurious correlations. Generate a color Ce uniformly at random. Select the pair (Xeg , C
e) with probability 1− ψe if the label Y eg and Ce are the same, else select them with a probability ψe. If the pair is selected, then color the image Xe ← T (Xg, Ce) and Y e ← Y eg . Selection probability ψe for the three environments are [ψ1, ψ2, ψ3] = [0.1, 0.2, 0.9]. Due to the selection bias, color is spuriously correlated, P(Ce = 1|Y e = 1) varies drastically [0.9, 0.8, 0.1]. We provide the graphical models for the CMNIST variants and computations of P(Ce|Y e), P(Y e|Xe) in the supplement.

heading: 4.1 RESULTS
text: We use the first two environments (e = 1, 2) to train and third environment (e = 3) to test. Other details of the training (models, hyperparameters, etc.) are in the supplement. For each of the above datasets, we run the experiments for different amounts of training data from 1000 to up to 60000 samples (10 trials for each data size). In Figure 2, we compare the models trained using IRM and ERM in terms of the classification error on the test environment e = 3 (a poor performance indicates model exploits the color) for varying number of train samples. We also provide the performance of the ideal hypothetical optimal invariant model. Observe that except for in the covariate shift setting where IRM and ERM are similar as seen in Figure 2a (as predicted from Proposition 4), IRM outperforms ERM in the remaining three datasets (as predicted from Proposition 5) as seen in Figure 2b-d. We further validate this claim through the regression experiments provided in the supplement. In CF-CMNIST, IRM achievs an error of 0.45, which is much better than error of ERM (0.7) but is marginally better than a random guess. This suggests that confounder induced spurious correlations are harder to mitigate and may need more samples than needed in anti-causal case (AC-CMNIST).

heading: 5 CONCLUSION
text: We presented a sample complexity analysis of IRM to answer the question: when is IRM better than ERM (and vice-versa)? For distribution shifts such as the covariate shifts, we proved that both IRM and ERM have similar sample complexity and arrive at the desired OOD solution asymptotically. For distribution shifts involving confounders and (or) anti-causal variables and polynomial generative models, we proved that IRM is guaranteed to achieve the desired OOD solution while ERM can be asymptotically biased. We proposed new variants of original colored MNIST dataset from Arjovsky et al. (2019), which are more comprehensive and better capture how spurious correlations occur in reality. To the best of our knowledge, we believe this to be the first work that provides a rigorous characterization of impact of factors such as model complexity, number of environments on the sample complexity and its distance from the OOD solution in distribution shifts that go beyond covariate shifts.

heading: 6 ACKNOWLEDGEMENTS
text: This work was supported in part by the Rensselaer-IBM AI Research Collaboration (part of the IBM AI Horizons Network).

heading: 7 SUPPLEMENT

heading: 7.1 DIALOGUE ON IRM CONTINUED
text: [In the original IRM paper by Arjovsky et al. (2019), we left two graduate students ERIC and IRMA strolling along the Parisian streets on a warm summer evening. A lot has transpired since then. ERIC is now holed up inside his room at Cité Universitaire and IRMA has returned to her parents’ home in Provence. The two are talking through a Zoom call.]
• ERIC: After reading Arjovsky et al. (2019) and Gulrajani & Lopez-Paz (2020), I was not sure of how to understand the settings where IRM is beneficial over ERM and vice-versa? From Gulrajani & Lopez-Paz (2020), I understood that ERM continues to be the state-ofthe-art if sufficient care is taken in performing model selection.
• IRMA: But after reading this manuscript I understand that there are settings where no matter how one selects the model, ERM is bound to be at a disadvantage.
• ERIC: What you say seems to contradict my understanding of Gulrajani & Lopez-Paz (2020).
• IRMA: Actually, they ...
[IRMA’s audio and video freeze for a minute, leaving ERIC to ponder this conundrum for a little bit on his own. IRMA stops her video and continues only with audio and the conversation is able to resume.]
• IRMA: Sorry, how much did you hear? • ERIC: You were just starting to explain why ERM can be at a disadvantage, but I didn’t
hear anything after that.
• IRMA: Okay, let me start my explanation again. There is no contradiction. The current manuscript says that in covariate shift settings, there maybe no clear winner. Many of the datasets considered in Gulrajani & Lopez-Paz (2020) are perhaps similar to covariate shift settings. Also, there could be another reason why IRM did not outperform ERM in Gulrajani & Lopez-Paz (2020) and I will come to it in a bit.
• ERIC: Yes, you are right! The datasets are human labeled images if I recall correctly. In these datasets it’s safe to assume P(Y e|Xe) is close to invariant across domains (as long as the cohort of humans in the different domains are not very different). What happens when P(Y e|Xe) varies a lot? • IRMA: Yes, when P(Y e|Xe) varies a lot, IRM can be at an advantage over ERM. In this
manuscript, I learned that when data generation involves confounders or anti-causal variables, it is possible to show that IRM methods converge to desirable OOD solutions with good convergence rates.
• ERIC: Interesting! I should read that. The experiment on colored MNIST by Arjovsky et al. (2019) contained anti-causal variables right? Is that the reason why IRM performed better than ERM?
• IRMA: Yes, you are right. In the current manuscript, the authors also develop other variants of colored MNIST (covariate shift based, confounder based, anti-causal based, and hybrid of confounder and anti-causal). IRM retains its advantage on all but the first dataset. I liked that IRM has advantage over ERM in confounded datasets as I believe many real datasets can have confounders generating spurious correlations.
• ERIC: Wow! These different variants of CMNIST sound exciting. Going back to the covariate shift case, do you think that we can perhaps come up with a finer criterion to say when IRM is better than ERM and vice-versa?
• IRMA: Yes, actually I have been thinking about this problem and would be happy to share my initial thoughts. Identifying a representation Φ∗ that leads to invariant conditional distributions is crucial to the success of IRM. A subtle factor that is implicitly assumed is representations obtained from multiple domains should overlap.
• ERIC: Can you clarify what you mean by overlap? • IRMA: Sure! Consider the dataset from domain 1, say photographs of birds, and domain 2,
say sketches of birds. Imagine I have access to the oracle representation Φ∗. Now I pass the data from domain 1 and 2 through Φ∗ to get the representations. If these representations from the two domains live in very different parts of the representation space, then we cannot hope that IRM will offer any advantage. This is the other explanation, which I was going to come to, why IRM did not outperform ERM in Gulrajani & Lopez-Paz (2020).
• ERIC: How about when there is a complete overlap? • IRMA: Yes, if there is a strong or complete overlap in the representations from the two
domains, then it is possible that IRM can help. In fact I believe in such settings, if the ERM models trained on the two domains disagree a lot, then that can be a strong indication towards the models heavily exploiting spurious correlations. In such cases, I believe IRM can again offer some advantage.
• ERIC: I will try to put these ideas down on paper and meet you again for a discussion. • IRMA: Great! I truly hope that it can be in person at the café in Palais-Royal where we first
started this conversation. So long for now!
[The students end their call.]

heading: 7.2 SUPPLEMENTARY MATERIALS FOR EXPERIMENTS
text: In this section, we cover the supplementary materials for the experiments. The code to reproduce the results presented in this work can be found at https://github.com/IBM/OoD.

heading: 7.2.1 CLASSIFICATION
text: We first describe the model and other training details.
Choice of HΦ and other training details We use the same architecture for both ERM and IRM. We choose the architecture that was used in Arjovsky et al. (2019): 2 layer MLP. The first two layer consists of 390 hidden nodes, and the output layer has two nodes (for the two classes). We use ReLU activation in each layer, a regularization weight of 0.0011 is used for each layer. We use a learning rate of 4.9e-4, batch size of 512 for both ERM and IRM. We use 1000 gradient steps for IRM. As was done in the original IRM work (Arjovsky et al., 2019), we use a threshold on steps (190) after which a large penalty is imposed for violating the IRM constraint. We use the train domain validation set procedure described in Gulrajani & Lopez-Paz (2020) to select the penalty value from the set {1e4, 3.3e4, 6.6e4, 1e5} (with 4:1 train-validation split). With the same learning rate, we observed ERM was slower at learning than IRM. To ensure ERM always converges we set the number of epochs to a very high value 100 (118k steps).

heading: A. Covariate shift based for CMNIST
text: We provide the generative model for CS-CMNIST below.
Y eg ← L(Xeg ), Ce ← Uniform({0, 1}) Ue ← Bernoulli ( (Ce ⊕ Y eg )ψe + ( 1− (Ce ⊕ Y eg ) )( 1− ψe) )) Xe ← T (Xg, Ce)|Ue = 1, Y e ← Y eg |Ue = 1
(11)
A.1. Compute P(Ce|Y e) and P(Y e|Xe). We compute P(Ce|Y e) and P(Y e|Xe) for the covariate shift based CMNIST described by equation 11. P(Ce|Y e) helps us understand how the spurious correlations vary across the environments. P(Y e|Xe) helps us understand if the covariate shift condition is satisfied or not. Compute the probability P(Ce|Y e = 1) = P(Ce|Y eg = 1, Ue = 1) as follows.
P(Ce = 1|Y eg = 1, Ue = 1) = P(Ce = 1, Y eg = 1|Ue = 1)
P(Ce = 1, Y eg = 1|Ue = 1) + P(Ce = 0, Y eg = 1|Ue = 1)
P(Ce = 1, Y eg = 1|Ue = 1) = P(Ce = 1, Y eg = 1, Ue = 1)∑ a,b P(Ce = a, Y eg = b, Ue = 1) = 1 2 (1− ψe) P(Ce = 0, Y eg = 1|Ue = 1) = P(Ce = 0, Y eg = 1, Ue = 1)∑ a,b P(Ce = a, Y eg = b, Ue = 1) = 1 2 ψe P(Ce = 1|Y eg = 1, Ue = 1) = (1− ψe)
(12)
From the above simplification we gather that P(Ce = 1|Y e = 1) is 0.9, 0.8 and 0.1 in the three environments. Define pl = P (Y eg = 1|Xeg ). In MNIST data it is reasonable to assume deterministic labeling, i.e. pl = 1 or pl = 0. From the way the data is constructed we can assume that T is invertible, i.e. from colored image Xe we can get back the grayscale image and the color Xeg , C
e. We now move to computing P(Y e|Xe). We assume Y eg = 1 in the simplifcation below.
P(Y eg = 1, Ce = 0, Ue = 1|Xeg ) = P(Y eg = 1|Xeg )P(Ce = 0)P(Ue = 1|Y eg = 1, Ce = 0) = 0.5plψe P(Y eg = 0, Ce = 0, Ue = 1|Xeg ) = P(Y eg = 0|Xeg )P(Ce = 0)P(Ue = 1|Y eg = 1, Ce = 0) = 0.5(1− pl)(1− ψe) P(Y e = 1|Xe) = P(Y eg = 1|Xeg , Ce = 0, Ue = 1) =
P(Y eg = 1, Ce = 0, Ue = 1|Xeg ) P(Y eg = 1, Ce = 0, Ue = 1|Xeg ) + P(Y eg = 0, Ce = 0, Ue = 1|Xeg )
P(Y e = 1|Xe) = plψe 2plψe + 1− pl − ψe
(13)
Observe that under standard assumption of deterministic labeling pl = 1, P(Y e = 1|Xe) remains invariant and for high values of pl it is relatively stable across the environments.
A.2 Graphical model for covariate shift based CMNIST. In Figure 3, we provide the graphical model for covariate shift based CMNIST described in equation 11.
A.3. Results with numerical values and standard errors. In Table 2, we provide the numerical values for the results showed in the Figure 2 a) along with the standard errors.

heading: B. Colored MNIST with anti-causal variables (AC-CMNIST)
text: B.1. Compute P(Ce|Y e) and P(Y e|Xe). We first compute P(Ce|Y e). P(Ce = 1|Y e = 1) = P(Y e ⊕Ne = 1|Y e = 1) = P(Ne = 0|Y e = 1) = βe. Therefore, P(Ce = 1|Y e = 1) is 0.9, 0.8 and 0.1 in environments 1, 2, and 3 respectively. Next, we compute P(Y e|Xe). We assume Y eg = 1 in the simplfication below. We also assume deteministic labeling. In the simplfication that follows we use β = 0.25.
P(Y e = 1, Ce = 0|Xeg ) = P(Y e = 1, Ce = 0|Xeg ) = P(Y e = 1|Y eg = 1)P(Ce = 0|Y e = 1) = 0.75βe P(Y e = 0, Ce = 0|Xeg ) = P(Y e = 1, Ce = 0|Xeg ) = P(Y e = 0|Y eg = 1)P(Ce = 0|Y e = 0) = 0.25(1− βe) P(Y e = 1|Xe) = P(Y e = 1|Xeg , Ce = 0) = 3βe
2βe + 1
P(Y e = 1|Xe) = 0.25 (For environment 1) P(Y e = 1|Xe) = 0.428 (For environment 2) P(Y e = 1|Xe) = 0.964 (For environment 2)
(14)
B.2 Graphical model for anti-causal CMNIST In Figure 4, we provide the graphical model for AC-CMNIST described in equation 10 (for G = 1).
B.3. Results with numerical values and standard errors. In Table 3, we provide the numerical values for the results showed in the Figure 2 c) along with the standard errors.

heading: C. Colored MNIST with confounded variables (CF-CMNIST)
text: C.1. Compute P(Ce|Y e) and P(Y e|Xe). We start by computing P(Ce|Y e). Recall that Ce = (N ⊕Ne). In the simplfication that follows we use β = 0.25.
P(N = 0|Y e = 1) = P(N = 0, Y e = 1)
P(N = 0, Y e = 1) + P(N = 1, Y e = 1) = 0.75 ∗ 0.5 0.75 ∗ 0.5 + 0.25 ∗ 0.5 = 0.75
(15)
We use the above to compute P(Ce = 0|Y e = 1) = P(N = 0, Ne = 0|Y e = 1) +P(N = 1, Ne = 1|Y e = 1) = (1−βe)0.75 + 0.25βe = 0.75− 0.5βe. For environment 1, 2, 3 the above probability P(Ce = 0|Y e = 1) is 0.7, 0.65 and 0.30 respectively. Next, we compute the probability P(Y e|Xe). Suppose Y eg = 1 for the calculation below. We also assume deterministic labeling.
P(Y e = 1, Ce = 1|Xeg ) = P(Y e = 1, Ce = 1|Xeg ) = P(Y e = 1|Y eg = 1)P(Ce = 1|N = 0) = 0.75βe P(Y e = 0, Ce = 1|Xeg ) = P(Y e = 1, Ce = 1|Xeg ) = P(Y e = 0|Y eg = 1)P(Ce = 1|N = 1) = 0.25(1− βe) P(Y e = 1|Xe) = P(Y e = 1|Xeg , Ce = 1) = 3βe
2βe + 1
P(Y e = 1|Xe) = 0.25 (For environment 1) P(Y e = 1|Xe) = 0.428 (For environment 2) P(Y e = 1|Xe) = 0.964 (For environment 2)
(16)
C.2 Graphical model for confounded CMNIST. In Figure 4, we provide the graphical model for confounded CMNIST described in equation 10 (for G = 0).
C.3. Results with numerical values and standard errors. In Table 4, we provide the numerical values for the results showed in the Figure 2 b) along with the standard errors.

heading: D. Colored MNIST with anti-causal variables and confounded variables (HB-CMNIST)
text: D.1. Compute P(Ce|Y e) and P(Y e|Xe). We start by computing P(Ce|Y e). Recall that Ce = G(Y e⊕Ne) + (1−G)(N ⊕Ne). G = 1 with probability θ and 0 otherwise.
P(Ce = 1|Y e = 1) = P(C e = 1, Y e = 1)
P(Ce = 1, Y e = 1) + P(Ce = 0, Y e = 1) (17)
P(Ce = 1, Y e = 1) = θ(1− βe) + (1− θ)(0.25 + 0.5βe) (18) P(Ce = 0, Y e = 1) = θ(βe) + (1− θ)(0.75− 0.25βe) (19)
We used θ = 0.8 in the experiments. P(Y e|Xe) can be computed on the same lines as was shown for anti-causal and confounded model and it varies significantly across the environments.
D.2 Graphical model for confounded CMNIST.
In Figure 4, we provide the graphical model for confounded CMNIST described in equation 10 (for 0 < θ < 1).
D.3. Results with numerical values and standard errors. In Table 5, we provide the numerical values for the results showed in the Figure 2 d) along with the standard errors.

heading: 7.2.2 REGRESSION
text: We use the same structure for the generative model as described by Arjovsky et al. (2019). We work with the four different variants on the same lines as CMNIST (covariate shift based, confounded, anti-causal, hybrid). The comparisons in Arjovsky et al. (2019) were for anti-causal and hybrid models. The general model is written as
He ← N (0, σ2e Is) Xe1 ← N (0, σ2e Is) +Wh→1He
Y e ←W e1→yXe1 +N (0, σ2e) +Wh→yHe Xe2 ←Wy→2Y e +N (0, Is) +Wh→2He
(20)
He is the hidden confounder, Xe = [Xe1 , X e 2 ] is the observed covariate vector, Y e is the label. Different W ’s correspond to the weight vectors that multiply with the covariates and the confounders. The four datasets differ in the weight W vectors and we describe them below. σe is environment
dependent standard deviation, we use two training environments with σ1 = 0.2 and σ2 = 2.0 respectively and both environments .
• Covariate shift case (CS-regresion): In this case, we fix Wh→2, Wh→y , and Wy→2 to zero and we draw each entry of W1→y from 1sN (0, 1) and set Wh→1 to identity.
• Confounded variable case (CF-regression): Set Wy→2 to zero and we draw each entry in W1→y,Wh→1,Wh→2,Wh→y from 1sN (0, 1).
• Anti-causal variable case (AC-regression): Set Wh→y ,Wh→1 and Wh→2 to zero and draw each entry of W1→y and Wy→2 from 1sN (0, 1).
• Hybrid confounded and anti-causal variable case (HB-regression): Draw each Wh→y ,Wh→1, Wh→2, W1→y and Wy→2 from 1sN (0, 1).
We also present the graphical models for the four types of models used in Figure 5.

heading: 7.2.3 CHOICE OF HΦ AND OTHER TRAINING DETAILS
text: We use a linear model that takes as input Xe. For ERM we use standard linear regression from sklearn. For IRM, we use 50k gradient steps with learning rate 1e-3, batch size is equal to the size of the training data. We use the train domain validation set procedure described by Gulrajani & Lopez-Paz (2020) to select the penalty value from the set {0, 1e− 5, 1e− 4, 1e− 3, 1e− 2, 1e− 1} (with 4:1 train-validation split). We average the results over 25 trials.

heading: 7.2.4 RESULTS
text: We discuss results for the case when the length of the covariate vector Xe is 10. The desired optimal invariant predictor is W ∗ = [W1→y, 0]. We will compare ERM and IRM in terms of the model estimation error, i.e., the distance between the model estimated by the method Ŵ and the true model given as ‖Ŵ − W ∗‖2. In Figures 6, 7, 8, 9, we compare the model estimation error vs. the number of samples when the model. In these comparisons, we see that consistent with the classification experiments and predictions from Proposition 4 in the covariate shift case (See Figure 6) there is no clear winner between the two approaches. There are gains from using IRM in the other cases (Figures 7, 8, 9). However, in the confounder case in Figure 7, the gains from IRM appear in the low sample regime but are not there in the high sample regime. This is because for this setup the asymptotic bias of ERM is also very small. In addition to the Figures 6, 7, 8, 9, we provide the tables (see Table 6, 7, 8, 9) with the numerical values for the mean model estimation (and the standard error) error shown in the figures.
Figure 7: Comparisons: n = 10 CF-regression
0 500 1000 1500 2000 2500 3000 3500 4000 Number of samples
0.00
0.25
0.50
0.75
1.00
1.25
1.50
1.75
M od
el e
st im
at io
n er
ro r ERM
IRM
Figure 8: Comparisons: n = 10 AC-regression 0 500 1000 1500 2000 2500 3000 3500 4000 Number of samples
0.00
0.25
0.50
0.75
1.00
1.25
1.50
1.75
M od
el e
st im
at io
n er
ro r ERM
IRM
Figure 9: Comparisons: n = 10 HB-regression

heading: 7.3 PROOFS FOR THE PROPOSITIONS
text: In the results to follow, we will rely on Hoeffding’s inequality. We restate the inequality below for convenience. Lemma 1. (Hoeffding’s inequality). Let θ1, . . . θm be a sequence of i.i.d. random variables and assume that for all i, E[θi] = µ and P[a ≤ θi ≤ b] = 1. Then, for any > 0
P [ 1
m ∣∣ m∑ i=1 θi − µ ∣∣ > ] ≤ 2 exp(−2 m 2 (b− a)2 )
We restate the propositions from the main body. Next, we prove Proposition 1 from the main body of the manuscript.
Proposition 7. If ` is square loss, and Assumptions 1,2 hold, then m ◦Φ∗ solves the OOD problem (equation 1).
Proof. Define a predictor w ◦Φ∗, where Φ∗ is defined in Assumption 1. Let us simplify the expression for the risk for this predictor Re(w ◦ Φ∗) using square loss for `. Recall Ze = Φ∗(Xe)
Re(w ◦ Φ∗) = Ee [( Y e − Ee [ Y e ∣∣Ze]+ Ee[Y e∣∣Ze]− (w ◦ Φ∗)(Xe))2]
= Ee [( Y e − Ee [ Y e ∣∣Ze])2]+ Ee[(Ee[Y e∣∣Ze]− (w ◦ Φ∗)(Xe))2]+
2Ee [( Y e − Ee [ Y e ∣∣Ze])(Ee[Y e∣∣Ze]− (w ◦ Φ∗)(Xe))]
= Ee [( Y e −m(Ze) )2] + Ee [( m(Ze)− w(Ze) )2] + 2Ee [( Y e −m(Ze) )( m(Ze)− w(Ze) )] = Ee [( Y e −m(Ze) )2] + Ee [( m(Ze)− w(Ze)
)2] = ξ2 + Ee [( m(Ze)− w(Ze)
)2] (21)
In the above simplification in equation 21, we use the following equation 22 and equation 23, which rely on the law of total expectation.
Ee [( Y e −m(Ze) )( m(Ze)− w(Ze) )] = Ee [ Ee [( Y e −m(Ze) )( m(Ze)− w(Ze) )∣∣∣Ze]] = Ee [( Ee [ Y e ∣∣Ze]−m(Ze))(m(Ze)− w(Ze))] = 0
(22)
Ee [( Y e −m(Ze) )2] = Ee [ Ee [( Y e −m(Ze) )2∣∣∣Ze]] = Ee[Vare[Y e∣∣Ze]] = ξ2 (23) In the last equality in equation 23, we use the Assumption 1 and obtain ξ2. Therefore, substituting w = m in equation 21 achieves a risk of ξ2 for all the environments.
∀e ∈ Eall, Re(m ◦ Φ∗) = ξ2 (24)
Consider the environment q that satisfies Assumption 2. Let us simplify the expression for the risk achieved by every predictor f ∈ F in the environment q following the steps similar to equation 21.
Rq(f) = Eq [( Y q − Eq[Y q|Zq] )2] +
Eq [( Eq[Y q|Zq]− f(Xq) )2] + 2Eq [( Y q − Eq[Y q|Zq] )( Eq[Y q|Zq]− f(Xq) )] = Eq [( Y q −m(Zq) )2] + Eq [( m(Zq)− f(Xq) )2] + 2Eq [( Y q −m(Zq) )( m(Zq)− f(Xq)
)] = Eq [( Y q −m(Zq) )2] + Eq [( m(Zq)− f(Xq)
)2] = ξ2 + Eq [( m(Zq)− f(Xq)
)2] (25)
In the above simplification in equation 25, we use the following equation 26 Eq [( Y q −m(Zq) )( m(Zq)− f(Xq)) )] = Eq [ Ee [( Y e −m(Zq) )( m(Zq)− f(Xq)
)∣∣∣Zq]] = Eq [( Eq [ Y q ∣∣Zq]−m(Zq))(m(Zq)− f(Xq))] = 0 (Last equality follows from the Assumption 2)
(26)
Therefore, for environment q satisfying Assumption 2 from equation 25, it follows that for all f ∈ F , Rq(f) ≥ ξ2. Therefore, we can write that
∀f ∈ F , max e∈Eall Re(f) ≥ Rq(f) ≥ ξ2 (27)
Therefore, from equation 27 it directly follows that
min f∈F max e∈Eall
Re(f) ≥ ξ2 (28)
We showed in equation 24, Re(m ◦ Φ∗) = ξ2 for all the environments. Hence, f = m ◦ Φ∗ achieves the RHS of equation 28. This completes the proof.
Some of the proofs that we describe next take a few intermediate steps to build. Here we give a brief preview of the key ingredients that we developed to build these propositions.
• In Proposition 2, our goal is to carry out a sample complexity analysis of EIRM in the same spirit as ERM. However, there are two key challenges that we are faced with - i) the IRM penalty R ′ is not separable (as it is composed of terms involving squares of expectations)
and ii) unlike ERM, IRM is a constrained optimization problem. To deal with i), we develop an estimator in the next section that allows us to re-express IRM penalty in a separable fashion. To deal with ii), we define a parameter κ that measures the minimum separation between IRM penalty and . We show that as long as this separation for all the predictors in the hypothesis class is positive, then we can rely on κ-representative property from ShalevShwartz & Ben-David (2014) applied to the new estimator that we build to show that the set of empirical invariant predictors Ŝ IV( ) are the same as exact invariant predictors S IV( ).
• In Proposition 5, our goal is to show that approximate OOD can be achieved by IRM in the finite sample regime. This result builds on the infinite sample result from Arjovsky et al. (2019). In Theorem 9 in Arjovsky et al. (2019), it was shown that for linear models (defined in Assumption 5) obeying linear general position 6, if the gradient constraints in the exact gradient constraint IRM equation 4 are satisfied, then the OOD solution is achieved. We extend this result to show that if the constraints in the approximate penalty based IRM in equation 5 are satisfied, then we are guaranteed to be in the √ neighborhood of the
OOD solution. Note that this result is again in the infinite sample regime as it proves the approximation for solutions of the problem equation 5, which involves expectations w.r.t true distributions. Next, we exploit similar tools that we introduced to prove Proposition 2 to also prove the finite sample extension.
• In later sections, we show the generalizations to infinite hypothesis classes. In particular, we focus on parametric model families that are Lipschitz continuous. The extension to infinite hypothesis classes is based on carefully exploiting the covering number based techniques Shalev-Shwartz & Ben-David (2014) for the IRM penalty estimator that we introduced. We also provide generalizations of the results for linear models to polynomial models. To arrive at these results, we exploit some standard properties of tensor products.
7.3.1 EMPIRICAL ESTIMATOR OF R ′
Next, we define an estimator for R ′ . We first simplify R ′ as follows.
Observe that
∇w|w=1.0Re(w · Φ) = ∂Ee
[ `(w · Φ(Xe), Y e) ] ∂w ∣∣∣ w=1.0 = Ee [ ∂`(w · Φ(Xe), Y e) ∂w ∣∣∣ w=1.0 ] and
‖∇w|w=1.0Re(w·Φ)‖2 = (∂Ee[`(w · Φ(Xe), Y e)]
∂w
∣∣∣ w=1.0 )2 = ( Ee [ ∂`(w · Φ(Xe), Y e) ∂w ∣∣∣ w=1.0 ])2 (29)
In the above simplification, we used Leibniz integral rule and take the derivative inside the expectation.
Also we can write E[X]2 = E[AB], where A and B are independent and identical random variables with same distribution as X . Therefore, we consider two independent data points (Xe, Y e) ∼ Pe and (X̃e, Ỹ e) ∼ Pe.
‖∇w|w=1.0Re(w · Φ)‖2 = Ee [(∂`(w · Φ(Xe), Y e)
∂w
∣∣∣ w=1.0 )(∂`(w · Φ(X̃e), Ỹ e) ∂w ∣∣∣ w=1.0 )] (30)
In the above the expectation Ee is taken over the joint distribution over pairs of distributions of pairs (Xe, Y e), (X̃e, Ỹ e) from the same environment e.
We write
R ′ (Φ) = ∑ e∈Etr πe‖∇w|w=1.0Re(w · Φ)‖2
= ∑ e∈Etr πeEe [(∂`(w · Φ(Xe), Y e) ∂w ∣∣∣ w=1.0 )(∂`(w · Φ(X̃e), Ỹ e) ∂w ∣∣∣ w=1.0 )] (31) In the above simplification, we used equation 30. Define a joint distribution P̃ over the tuple (e, (Xe, Y e), (X̃e, Ỹ e)), where e ∼ {πo}o∈Etr , (Xe, Y e) ∼ Pe and (X̃e, Ỹ e) ∼ Pe. Also,
P̃ ( (e, (Xe, Y e), (X̃e, Ỹ e)) ) = πePe(Xe, Y e)Pe(X̃e, Ỹ e) (32)
We rewrite the above expression equation 31 in terms of an expectation w.r.t P̃, which we represent as Ẽ follows
R ′ (Φ) = Ẽ [(∂`(w · Φ(Xe), Y e) ∂w ∣∣∣ w=1.0 )(∂`(w · Φ(X̃e), Ỹ e) ∂w ∣∣∣ w=1.0 )] (33)
Define
` ′ ( h, ( (Xe, Y e), (X̃e, Ỹ e) )) = (∂`(w · Φ(Xe), Y e)
∂w
∣∣∣ w=1.0 )(∂`(w · Φ(X̃e), Ỹ e) ∂w ∣∣∣ w=1.0 ) (34)
Substitute equation 34 in equation 33 to obtain
R ′ (Φ) = Ẽ [ ` ′ ( h, ( (Xe, Y e), (X̃e, Ỹ e) ))] (35)
We construct a simple estimator R̂ ′ (Φ) by pairing the data points in each environment. For simplicity assume that each environment has even number of points. In environment e, which has ne points we construct ne2 pairs. Define a set of such pairs as
D̃ = {{(xe2i−1, ye2i−1), (xe2i, ye2i)} ne 2 i=1}e∈Etr (36)
R̂ ′ (Φ) =
2 |D| ∑ e∈Etr ne 2∑ i=1 ` ′ ( h, ( (Xe2i−1, Y e 2i−1), (X̃ e 2i, Ỹ e 2i) ))
(37)
There can be other estimators of R ′ (Φ) where we separately estimate each term ‖∇w|w=1.0Re(w.Φ)‖2 and πe in the summation. We rely on the above estimator equation 37 as its separability allows us to use standard concentration inequalities, e.g., Hoeffding’s inequality.
7.3.2 -REPRESENTATIVE TRAINING SET FOR R AND R ′
We use the definition of -representative sample from Shalev-Shwartz & Ben-David (2014) and state it appropriately for both R and R ′ .
Definition 1. A training set S is called -representative (w.r.t. domain Z , hypothesis H, loss ` and distribution D) if
∀h ∈ H, |R̂(h)−R(h)| ≤ (38) where R(h) = ED[`(h(X), Y ] and (X,Y ) ∼ D.
Following the above definition, we apply it to the set of points D̃ defined above in equation 36. D̃ is called -representative w.r.t. domain X , hypothesis HΦ, loss ` ′ (equation 33) and distribution P̃ (equation 32) if ∀Φ ∈ HΦ, |R̂ ′ (Φ)−R ′ (Φ)| ≤ (39)
where R ′ (Φ) = Ẽ [ ` ′ ( h, ( (Xe, Y e), (X̃e, Ỹ e) ))] (from equation 34) and
(e, (Xe, Y e), (X̃e, Ỹ e)) ∼ P̃.
Recall the definition of κ, κ = minΦ∈HΦ |R ′ (Φ) − |. Next, we show that if D̃ is κ2 -representative w.r.t X ,HΦ, loss ` ′ , distribution P̃ then the set of invariant predictors in equation 6 (Ŝ IV( ) ) and the set of invariant predictors in equation 5 (S IV( )) are equal.
Lemma 2. If κ > 0 and D̃ is κ2 -representative w.r.t X , HΦ, loss ` ′
and distribution P̃ , then Ŝ IV( ) = S IV( ).
Proof. First we show S IV( ) ⊆ Ŝ IV( ). From the definition of κ, κ = minΦ∈HΦ |R ′ (Φ) − | it follows that ∀Φ ∈ HΦ
|R ′ (Φ)− | ≥ κ =⇒ R ′ (Φ) ≥ + κ or R ′ (Φ) ≤ − κ (40)
Consider any Φ in S IV( ). R ′ (Φ) ≤ (41)
Given the definition of κ and equation 40, we obtain
R ′ (Φ) ≤ =⇒ R ′ (Φ) ≤ − κ (42)
Therefore, S IV( ) ⊆ S IV( − κ). Also, it follows from the definition of the set S IV( ) that S IV( − κ) ⊆ S IV( ). Hence,
S IV( ) = S IV( − κ) (43)
Consider any Φ in S IV( )
R ′ (Φ) ≤ − κ (From equation 43)
R ′ (Φ)− R̂ ′ (Φ) + R̂ ′ (Φ) ≤ − κ
R̂ ′ (Φ) ≤ − κ+ |R ′ (Φ)− R̂ ′ (Φ)|
(44)
From the definition of κ2 -representativeness it follows that |R ′ (Φ) − R̂′(Φ)| ≤ κ2 and substituting this in equation 44 we get
R̂ ′ (Φ) ≤ − κ/2 =⇒ R̂ ′ (Φ) ≤ =⇒ S IV( ) ⊆ Ŝ IV( ) (45)
Next we show Ŝ IV( ) ⊆ S IV( ).
Consider Φ ∈ Ŝ IV( )
R̂ ′ (Φ) ≤
R̂ ′ (Φ)−R ′ (Φ) +R ′ (Φ ′ ) ≤
R ′ (Φ) ≤ + |R̂ ′ (Φ)−R ′ (Φ)|
(46)
From the definition of κ2 -representativeness it follows that |R ′ (Φ) − R̂′(Φ)| ≤ κ2 and substituting this in equation 46 we get
R ′ (Φ) ≤ + κ
2 (47)
From equation 40, it follows that R ′ (Φ) ≤ + κ2 =⇒ R ′ (Φ) ≤ . Therefore, Φ ∈ S IV( ). This proves the second part Ŝ IV( ) ⊆ S IV( ) and completes the proof.
Lemma 3. If κ > 0, D is ν2 -representative w.r.t X ,HΦ, loss ` and distribution P̄ (joint distribution over (e,Xe, Y e) defined in Section 3.1) and and D̃ is κ2 -representative w.r.t X , HΦ, loss ` ′ and distribution P̃, then every solution Φ̂ to EIRM (equation 6) satisfies Φ̂ is in S IV( ) and R(Φ∗) ≤ R(Φ̂) ≤ R(Φ∗) + , where Φ∗ is the solution of IRM in equation 5.
Proof. Given the condition in the above lemma, we are able to use the previous Lemma 2 to deduce that Ŝ IV( ) = S IV( ). This makes the set of predictors satisfying the constraints in EIRM equation 6 and IRM equation 5 the same.
Φ∗ solves equation 5 and Φ̂ solves equation 6. From ν2 -representativeness we know thatR(Φ̂)− ν 2 ≤ R̂(Φ̂). From the optimality of Φ̂ we know that R̂(Φ̂) ≤ R̂(Φ∗) (Φ∗ ∈ Ŝ IV( ) = S IV( )). Moreover, from ν2 -representativeness we know that R̂(Φ
∗) ≤ R(Φ∗) + ν2 . We combine these conditions as follows.
R(Φ̂)− ν 2 ≤ R̂(Φ̂) ≤ R̂(Φ∗) ≤ R(Φ∗) + ν 2 (48)
Comparing the first and third inequality in the above equations we getR(Φ̂) ≤ R(Φ∗)+ν. From the optimality of Φ∗ over the set S IV( ) and since Φ̂ ∈ S IV( ) it follows that R(Φ∗) ≤ R(Φ̂). Hence, R(Φ∗) ≤ R(Φ̂) ≤ R(Φ∗) + ν. This completes the proof.
Next, we prove Proposition 2 from the main body of the manuscript.
Proposition 8. For every ν > 0, > 0 and δ ∈ (0, 1), ifHΦ is a finite hypothesis class, Assumption 3 holds, κ > 0, and if the number of samples |D| is greater than max { 16L′4
κ2 , 8L2 ν2
} log ( 4|HΦ|
δ
) ,
then with a probability at least 1 − δ, every solution Φ̂ of EIRM (equation 6) is a ν approximation of IRM, i.e. Φ̂ ∈ S IV( ), R(Φ∗) ≤ R(Φ̂) ≤ R(Φ∗) + ν, where Φ∗ is a solution to IRM (equation 5).
Proof. From Lemma 3, we know that if D is ν2 -representative w.r.t X , HΦ, loss ` and distribution P and if D̃ is κ2 -representative X , HΦ, loss ` ′ and distribution P̃, then the the claim in the above Proposition is true.
Define an event A: D is ν2 -representative w.r.t X ,HΦ, loss ` and distribution P̄.
Define an event B: D̃ is κ2 -representative X ,HΦ, loss ` ′ and distribution P̃.
Define success as A ∩ B. Next, we show that if |D| is greater than max { 16L′4
κ2 , 8L2 ν2
} log ( 4|HΦ|
δ
) ,
then P(A∩B) occurs with a probability at least 1−δ. P (A∩B) = 1−P (Ac∪Bc) ≥ 1−P (Ac)−
P (Bc). If we bound P (Ac) ≤ δ2 and P (B c) ≤ δ2 , then we know the probability of success is at least 1− δ. We write
P(A) = P ({ D : ∀h ∈ HΦ, |R̂(h)−R(h)| ≤ ν
2
}) = 1−P ({ D : ∃h ∈ HΦ, |R̂(h)−R(h)| > ν
2
})
P ({ D : ∃h ∈ HΦ, |R̂(h)−R(h)| > ν
2
}) = P ( ⋃ h∈HΦ { D : |R̂(h)−R(h)| > ν 2 }) ≤ ∑ h∈HΦ P ({ D : |R̂(h)−R(h)| > ν 2 }) (49)
The loss function is bounded |`(Φ(·), ·)| ≤ L. From Hoeffding’s inequality in Lemma 1 it follows that
P ({ D : |R̂(h)−R(h)| > ν
2
}) ≤ 2 exp ( − |D|ν 2
8L2
) (50)
Using this expression equation 50 in equation 49, we get
2|HΦ| exp ( − |D|ν 2
8L2
) ≤ δ
2 =⇒ |D| ≥ 8L
2 ν2 log (4|HΦ| δ ) (51)
P(B) = P ({ D̃ : ∀h ∈ HΦ, |R̂ ′ (h)−R ′ (h)| ≤ κ
2
}) = 1−P ({ D̃ : ∃h ∈ HΦ, |R̂ ′ (h)−R ′ (h)| > κ
2
})
P ({ D̃ : ∃h ∈ HΦ, |R̂ ′ (h)−R ′ (h)| > κ
2
}) = P ( ⋃ h∈HΦ { D̃ : |R̂ ′ (h)−R ′ (h)| > κ 2 }) ≤ ∑ h∈HΦ P ({ D̃ : |R̂ ′ (h)−R ′ (h)| > κ 2 }) (52)
The gradient of loss function is bounded |∂`(h(·),·)∂w |w=1.0| ≤ L ′. From the definition of ` ′ (h(·), ·) in equation 34, we can infer that |`′(h(·), ·)| ≤ L′2. Recall that R′(h) =
Ẽ [ ` ′ ( h, ( (Xe, Y e), (X̃e, Ỹ e) ))] From Hoeffding’s inequality in Lemma 1 it follows that
P ({ D̃ : |R̂ ′ (h)−R ′ (h)| > κ
2
}) ≤ 2 exp ( − |D̃|κ 2
8L′4
) = 2 exp ( − |D|κ 2
16L′4
) (53)
Using the above equation 53 in equation 52 we get
2|HΦ| exp ( − |D|κ 2
16L′4
) ≤ δ
2 =⇒ |D| ≥ 16L
′4 κ2 log (4|HΦ| δ ) (54)
Combining the two conditions in equation 51 and equation 51 we get that if
|D| ≥ max {16L′4
κ2 ,
8L2
ν2
} log (4|HΦ|
δ ) then with probability at least 1− δ event A ∩B occurs.

heading: 7.3.3 PROPERTY OF LEAST SQUARES OPTIMAL SOLUTIONS
text: We first remind ourselves of a simple property of least squares minimization. Consider the least squares minimization setting, where R(h) = E[(Y − h(X))2].
E[(Y − h(X))2] = E [( Y − E [ Y ∣∣X]+ E[Y ∣∣X]− h(X))2]
= EX [ E [( Y − E [ Y ∣∣X])2|X]]+ EX[(E[Y ∣∣X]− h(X))2]
= EX [ Var [ Y ∣∣X]]+ EX[(E[Y ∣∣X]− h(X))2]
(55)
In the above simplification, we use the law of total expectation. Both the terms in the above equations are always greater than or equal to zero. The first term does not depend on h, which implies the minimization can focus on second term only.
min h R(h) = EX
[ Var [ Y ∣∣X]]+ min h EX [( E [ Y ∣∣X]− h(X))2] (56)
Assume that P has full support over X . Define ∀x ∈ X , h∗(x) = E[Y |X = x]. ∀h,R(h) ≥ EX [ Var [ Y ∣∣X]]. Since R(h∗) = EX[Var[Y ∣∣X]]. Therefore,
h∗ ∈ arg min h
EX [( E [ Y ∣∣X]− h(X))2] (57)
Moreover, we conclude that h∗ is the unique minimizer. Observe that EX [( E [ Y ∣∣X]− h(X))2] is zero for h = h∗. From Theorem 1.6.6 in (Ash et al., 2000), it follows that any other minimizer is same as h∗ except over a set of measure zero.
Recall the definition of m(x) = E[Y e|Xe = x] Lemma 4. Let ` be the square loss. If Assumption 4 holds and m ∈ HΦ, then m uniquely solves expected risk minimization m ∈ arg minΦ∈HΦ R(Φ) and also uniquely solves IRM (equation 5).
Proof. R(Φ) = ∑ e∈Etr π
eRe(Φ). From Assumption 4 and the observation in equation 57, it follows that the unique optimal solution to expected risk minimization for each Re is m. Therefore, m also minimizes the weighted combination R.
To show the latter part of the Lemma, if we can show that m ∈ S IV( ), then the rest of the proof follows from the previous part as we already showed m is a minimizer among all the functions in HΦ and S IV( ) ⊆ HΦ. Supposem 6∈ S IV( ). This implies there exists at least one environment for which ‖∇w|w=1.0Re(w · m)‖2 > 0 =⇒ ∇w|w=1.0Re(w ·m) 6= 0. As a result ∃ w in the neighborhood of w = 1.0 where Re(w ·m) < Re(m) (if such a point does not exist and all the points in the neigbohood of w = 1.0 are greater than or equal to Re(m) that would make ∇w|w=1.0Re(w · m) = 0, which would be a contradiction). Therefore, Re(w · m) < Re(m). However, this is a contradiction as we know that m is the unique optimizer for each environment. Hence, m 6∈ S IV( ) cannot be true and thus m ∈ S IV( ). This completes the proof.
Next, we prove Proposition 4 from the main body of the manuscript. Proposition 9. Let ` be the square loss. For every ν > 0, > 0 and δ ∈ (0, 1), if HΦ is a finite hypothesis class, m ∈ HΦ, Assumptions 3, 4 hold, and
• if the number of samples |D| is greater than max { 8L2
ν2 log( 4|HΦ| δ ), 16L′4 2 log( 2 δ ) }
, then with a probability at least 1−δ, every solution Φ̂ to EIRM (equation 6) satisfiesR(m) ≤ R(Φ̂) ≤ R(m)+ ν. If also ν < κ̃, then Φ† = m.
• if the number of samples |D| is greater than 8L 2 ν2 log( 2|HΦ| δ ), then with a probability at least 1− δ, every solution Φ† to ERM satisfies R(m) ≤ R(Φ†) ≤ R(m) + ν. If also ν < κ̃, then Φ† = m.
Proof. We first cover the second part of the Proposition. From Proposition 3, we know that the output of ERM will satisfy R(Φ+) ≤ R(Φ†) ≤ R(Φ+) + ν. In this case from Lemma 4, it follows that Φ+ = m. From the definition of κ̃ and the fact that ν < κ̃ implies that Φ† = m. We now move to the first part of the Proposition.
For EIRM we will derive a tighter bound on sample complexity than the one in Proposition 2 since we can now use the Assumption 4. Observe that ∀e ∈ Etr, ∇w|w=1.0Re(w ·m) = 0 (see the proof of Lemma 4). Therefore, R ′ (m) = 0.
Define an event A: D is ν2 -representative w.r.t X ,HΦ, loss ` and distribution P̄.
Define an event B: D̃ is such that |R̂′(m)− R′(m)| ≤ 2 . Since R ′ (m) = 0, |R̂′(m)− R′(m)| ≤ 2 =⇒ |R̂ ′ (m)| ≤ 2 =⇒ R̂ ′ (m) ≤ 2 . Therefore, m ∈ Ŝ IV( ).
If A ∩ B occurs, then R(m) ≤ R(Φ̂) ≤ R(m) + ν; we justify claim next. Suppose Φ̂ solves equation 6. If eventA occurs, then from ν2 -representative condition we know thatR(Φ̂)− ν 2 ≤ R̂(Φ̂). From optimality of Φ̂ it follows that R̂(Φ̂) ≤ R̂(m) (event B =⇒ m ∈ Ŝ IV( )). Moreover, from ν 2 -representative property, we conclude that R̂(m) ≤ R(m) + ν 2 . We combine these conditions as follows.
R(Φ̂)− ν 2 ≤ R̂(Φ̂) ≤ R̂(m) ≤ R(m) + ν 2 (58)
From the above we have R(m) ≤ R(Φ̂) ≤ R(m) + ν. Recall the definition of κ̃ and since ν < κ̃ =⇒ Φ̂ = m. Next, we bound the probability of success. P (A ∩B) = 1− P (Ac ∪Bc) ≥ 1− P (Ac)− P (Bc). If we can bound P (Ac) ≤ δ2 and P (B
c) ≤ δ2 , then we know the probability of success is at least 1− δ. We write
P(A) = P ({ D : ∀h ∈ HΦ, |R̂(h)−R(h)| ≤ ν
2
}) = 1−P ({ D : ∃h ∈ HΦ, |R̂(h)−R(h)| > ν
2 }) From equation 51 if the condition
|D| ≥ 8L 2 ν2 log (4|HΦ| δ ) (59)
is true, then is true, then event Ac occurs with probability at most δ2 . We write P(B) = P ({ D̃ :, |R̂′(m) − R′(m)| ≤ 2 }) = 1 − P ({ D̃ : |R̂′(m) − R′(m)| > 2 }) . The gradient of loss function is bounded |∂`(Φ(·),·)∂w |w=1.0| ≤ L ′ . From Hoeffding’s inequality in Lemma 1 it follows that
P ({ D̃ : |R̂ ′ (h)−R ′ (h)| >
2
} ) ≤ 2 exp(− |D| 2
16L′4
) (60)
We bound the above equation 60 by δ2 to get
2 exp(− |D| 2 16L′4 ) ≤ δ 2 =⇒ |D| ≥ 16L
′4
2 log(
4 δ ) (61)
Combining the two conditions equation 59 and equation 61,
|D| ≥ max {8L2 ν2 log( 4|HΦ| δ ), 16L′4 2 log( 4 δ ) }
This ensures P (A ∩B) ≥ 1− δ. This completes the proof.
Before stating the proof of Proposition 5, we will prove an intermediate proposition. For clarity, we will restate the result (Theorem 9 from Arjovsky et al. (2019)) next.
Proposition 10. (Theorem 9 Arjovsky et al. (2019)) If Assumptions 5 and 6 (with r = 1) hold and let Φ ∈ Rn×1 (Φ 6= 0), then
ΦTEe[XeXe,T]Φ = ΦTEe[XeY e] (62)
holds for all e ∈ Etr iff Φ = S̃Tγ.
Next, we propose an -approximation of Proposition 10.
Define 0 = π min
|Etr| (ωλmin) 2(12− 8
√ 2).
Proposition 11. Let ` be the square loss. If Assumptions 5, 6 (with r = 1) and 7 hold , then for all 0 < < 0, the solution Φ
R ′ (Φ) ≤ (63)
satisfies Φ = S̃Tγ(α), where α ∈ [ 1 1+ 12ωλmin √ |Etr| πmin , 1 1− 12ωλmin √ |Etr| πmin ]
Proof. Let us start by simplifying ∇w|w=1.0Re(w · Φ), using square loss for ` and linear representation Φ ∈ Rn×1.
∇w|w=1.0Re(w · Φ) = ∂Ee
[( Y e − w · ΦTXe )2] ∂w ∣∣∣ w=1.0 = 2Ee [ (ΦTXe)2 ] − 2Ee [ ΦTXeY e ] = 2ΦTEe [ XeXe,T ] Φ− 2ΦTEe [ Y eXe
] (64)
Plug the above equation 64 in the condition R ′ (Φ) ≤ to get
R ′ (Φ) = ∑ e πe‖∇w|w=1.0Re(w · Φ)‖2 = 4 ∑ e πe ( ΦTEe [ XeXe,T ] Φ− ΦTEe [ XeY e ])2 ≤
(65)
From the bound on πe in Assumption 5 it follows that for each e ∈ Etr( ΦTEe [ XeXe,T ] Φ− ΦTEe [ XeY e ])2 ≤ |Etr|
4πmin∣∣∣ΦTEe[XeXe,T]Φ− ΦTEe[XeY e]∣∣∣ ≤√ |Etr| 4πmin
(66)
Note that if the condition above equation 66 is not true, then the preceeding condition in equation 65 cannot be true as contribution from one term in the summation itself will exceed . In the above equation 66, we are using the positive square root since RHS has to be greater than or equal to zero.
We compute the second derivative of loss w.r.t w
∇2wRe(w · Φ) = ∂2Ee
[( Y e − ΦTXew )2] ∂w2 = ∂ ( 2wEe [ (ΦTXe)2 ] − 2Ee [ ΦTXeY e ]) ∂w
= 2ΦTEe [ XeXe,T ] Φ
= 2ΦTΣeΦ
(67)
Since Σe is symmetric, we can use the eigenvalue decomposition of Σe = UΛUT in equation 67 to get ΦTΣeΦ = ΦTUΛUTΦ. Substitute Φ̃ = UTΦ to get ΦTΣeΦ = Φ̃TΛΦ̃ ≥ λmin(Σe)‖Φ̃‖2 = λmin(Σe)Φ
TUUTΦ = λmin(Σe)‖Φ‖2. From Assumption 5, we have λmin(Σe) ≥ λmin and from Assumption 7 we have ‖Φ‖2 ≥ ω. Therefore, we can deduce that 2ΦTΣeΦ ≥ 2λminω. Therefore, the second derivative defined in equation 67 is always greater than or equal to 2λminω.
∇2wRe(w · Φ) ≥ 2λminω > 0 (68)
Let ′ = 2 √
|Etr| 4πmin
= √
|Etr| πmin . We rewrite equation 66 in terms of∇w|w=1.0Re(w ·Φ) and ′ to get
|∇w|w=1.0Re(w · Φ)| ≤ ′
− ′ ≤ ∇w|w=1.0Re(w · Φ) ≤
′ (69)
Since the second derivative (equation 68) is strictly positive and larger than or equal to λminω =⇒ ∃ we in the neighborhood of w = 1.0 at which ∇w|w=weRe(w · Φ) = 0. This holds for all the environments in Etr. Define
c(w) = ∇wRe(w · Φ)
Also, define c ′ (w) = ∂c(w)∂w . Also, c ′ (w) = ∇2wRe(w·Φ). Suppose∇w|w=1.0Re(w·Φ) = c(1) < 0. Since for all w, c′(w) > 0 (from equation 68), ∃ we > 1, where c(we) = 0. Using fundamental theorem of calculus, we write
c(w)− c(1) = ∫ w
1
c ′ (u)du ≥ 2λminω(w − 1)
Substituting w = we in the above
c(we)− c(1) ≥ 2λminω(we − 1) −c(1) ≥ 2λminω(we − 1)
we ≤ 1− c(1) 2λminω
≤ 1 + ′
2λminω = 1 +
√ |Etr| πmin 2λminω (70)
Suppose ∇w,w=1.0Re(w.Φ) = c(1) > 0. Using fundamental theorem of calculus we can write
c(1)− c(w) = ∫ 1 w c ′ (u)du ≥ λminω(1− w)
Substituting w = we in the above
c(1)− c(we) ≥ 2λminω(1− we) c(1) ≥ 2λminω(1− we)
we ≥ 1− c(1) 2λminω
≥ 1− ′
2λminω = 1−
√ |Etr| πmin 2λminω (71)
Define η1 = 1 2ωλmin
√ |Etr| πmin
1− 12ωλmin √ |Etr| πmin , η2 = − 12ωλmin
√ |Etr| πmin
1+ 12ωλmin √ |Etr| πmin . Combining equation 70 and equation 71 and
using the definition of η1 and η2, we can conclude that we ∈ [ 11+η1 , 1 1+η2 ]. If we reparamterize we = 11+ηe , then η e ∈ [η2, η1]. We expand this condition∇w|w=weRe(w · Φ) = 0.
∂ ∂w Ee [( Y e − ΦTXew )2] = 2wEe [ (ΦTXe)2 ] − 2Ee [ ΦTXeY e ] = 2ΦT [ Ee [ XeXe,T ] Φwe − Ee [ XeY e
]] = 2ΦT [ Ee [ XeXe,T ] Φwe − Ee [ XeXe,T ] S̃Tγ − Ee [ Xeεe
]] = 2ΦT [ Ee [ XeXe,T ]( Φwe − S̃Tγ ) − Ee [ Xeεe
]] = 2ΦT [ Ee [ XeXe,T ]( Φwe − S̃Tγ ) − Ee [ Xeεe ]] = 0
= 2ΦT [ Ee [ XeXe,T ]( Φ 1
1 + ηe − S̃Tγ
) − Ee [ Xeεe ]] = 0
(72)
Assume that Φ 11+ηe 6= S̃ Tγ, for all e ∈ Etr. If this assumption is not true and Φ 11+ηe = S̃ Tγ for some e ∈ Etr, then it already establishes the claim we set out to prove in this Proposition (since ηe ∈ [η2, η1]).
Define qe = [ Ee [ XeXe,T ]( Φ 11+ηe − S̃ Tγ ) − Ee [ Xeεe ]] . From the Assumption 6 and since
Φ 11+ηe 6= S̃ Tγ, we know dim
( span { qe }) > n− 1.
From rank-nullity theorem we know dimension of kernel space of Φ (rank of Φ is 1) is n− 1. From equation 72 it follows that qe is in kernel space of Φ. Therefore, dim(Ker(Φ)) = n − 1 =⇒ dim ( span { qe }) ≤ n− 1 which leads to a contradiction.
Therefore, Φ 11+ηe = S̃ Tγ at least for one environment.

heading: If Φ = S̃Tγ(1 + ηe), then ‖Φ‖2 = ‖S̃Tγ‖2(1 + ηe)2 ≥ ‖S̃Tγ‖2(1 + η2)2 ≥ ‖S̃Tγ‖2 12 ≥
text: ω. In this simplification, we use (1 + η2)2 ≥ 1(
1+ 12ωλmin √ |Etr| πmin )2 ≥ 1( 1+ 12ωλmin √ 0|Etr| πmin )2 = 1(
1+ √ 3−2 √ 2 )2 = 12 and ‖S̃Tγ‖2 ≥ 2ω (from Assumption 7). This ensures that for any solution of the form Φ 11+ηe the assumption ‖Φ‖ 2 ≥ ω is automatically satisfied.
If Φ = S̃Tγ(1+ηe), then ‖Φ‖2 = ‖S̃Tγ‖2(1+ηe)2 ≤ ‖S̃Tγ‖2(1+η1)2 ≤ ‖S̃Tγ‖2 ( 3+2 √ 2 2 ) ≤ Ω.
In the above simplification, we use (1 + η1)2 ≤ 1( 1− 12ωλmin √ |Etr| πmin )2 ≤ 1( 1− 12ωλmin √ 0|Etr| πmin )2 = 1(
1− √ 3−2 √ 2 )2 = 3+2√22 and ‖S̃Tγ‖2 ≤ 23+2√2Ω (from Assumption 7). This ensures that for any solution of the form Φ 11+ηe the assumption ‖Φ‖ 2 ≤ Ω is automatically satisfied.
The entire proof so far has characterized the property of Φ that satisfies R ′ (Φ) ≤ . But how do we
know such a Φ exists. Recall S̃tγ ∈ HΦ. For each environment e ∈ Etr
∇w|w=1.0Re(w · Φ) = ΦTEe [ XeXe,T ] Φ− ΦTEe [ XeY e ] = γTS̃Ee [ XeXe,T ] S̃Tγ − γTS̃Ee [ XeY e
] = γTEe [ Ze1Z e,T 1 ] γ − γTEe [ Ze1Y e ]
= γTEe [ Ze1Z e,T 1 ] γ − γTEe [ Ze1Z e,T 1 ] γ − γTEe [ Ze1ε e ]
= 0
(73)
We use Assumption 5 in the simplification above in equation 73 Therefore, R ′ (S̃tγ) = 0 and as a result the existence of a Φ that satisfies R ′ (Φ) ≤ is guaranteed. This completes the proof.
Before proving Proposition 5, we first establish that Assumptions 5 and 7 are sufficient to ensure that Assumption 3 holds and we can thus use bounds L and L′ defined in Assumption 3.
Proving Assumption 3 for square loss ` from Assumptions 5, 7. From Assumption 5, we have
Ze = (Ze1 , Z e 2) Xe = SZe
‖Xe‖ ≤ ‖S‖‖Ze‖ (74)
From Assumption 5, ‖S‖ is bounded and ‖Ze‖ is bounded =⇒ ‖Xe‖ is bounded as well (from equation 74). Therefore, there exists Xsup <∞, s.t. ‖Xe‖ ≤ Xsup.
Y e = (S̃Tγ)TXe + εe
|Y e| ≤ ‖S̃Tγ‖‖Xe‖+ |εe| =⇒
|Y e| ≤
√ 2
3 + 2 √ 2 ΩXsup + εsup
(75)
In the last step of equation 75, we use ‖S̃Tγ‖2 ≤ 2 3+2 √ 2 Ω (Assumption 7), ‖Xe‖ ≤ Xsup derived above (equation 74), and |εe| ≤ εsup (Assumption 5). Therefore, Y e is bounded and there exists a K such that |Y e| ≤ K ≤ √ ΩXsup + εsup. Therefore, ∀Φ ∈ HΦ and for all Xe, Y e sampled from the model in Assumption 5 we have
`(Φ(Xe), Y e) = (Y e − ΦTXe)2 ≤ (K + √ ΩXsup)2 (76)
∣∣∣∂`(w · ΦTX,Y ) ∂w ∣∣∣ w=1.0 ∣∣∣ = ∣∣ΦTX(ΦTX − Y )∣∣ ≤ (√ΩXsup)(√ΩXsup +K) (77) From equation 76, we conclude that `(Φ(·), ·) is bounded and there exists an L such that |`(Φ(·), ·)| ≤ L ≤ (K + √ ΩXsup)2. From equation 77, we conclude that ∂`(w.Φ
TX,Y ) ∂w ∣∣∣ w=1.0
is bounded and there exists an L′ such that ∣∣∣ `(w.Φ(·),·)∂w ∣∣∣
w=1.0 ∣∣∣ ≤ L′ ≤ (√ΩXsup)(√ΩXsup +K). Define th = 24−16 √ 2 3 πmin |Etr| (ωλmin) 2 and τ = 12ωλmin √ 3|Etr| 2πmin
Next, we prove Proposition 5 from the main body of the manuscript.
Proposition 12. Let ` be the square loss. For every ∈ (0, th) and δ ∈ (0, 1), if Assumptions 5, 6 (with r = 1), 7 hold and if the number of data points |D| is greater than 16L ′4 2 log ( 2|HΦ|
δ
) , then with
a probability at least 1 − δ, every solution Φ̂ to EIRM (equation 6) satisfies Φ̂ = (S̃Tγ ) α, where
α ∈ [ 1 1+τ √ , 1 1−τ √ ].
Proof. Define an event A: {D̃ : ∀Φ ∈ HΦ, |R̂ ′ (Φ)−R′(Φ)| ≤ 2}. If event A happens, then
R̂ ′ (Φ) ≤
R̂ ′ (Φ)−R ′ (Φ) +R ′ (Φ) ≤
R ′ (Φ) ≤ + |R̂ ′ (Φ)−R ′ (Φ)|
R ′ (Φ) ≤ 3
2
(78)
If event A happens, then every solution Φ̂ to EIRM (equation 6) satisfies equation 78. From equation 78, we see that we should substitute with 3 2 and use Proposition 11. If the Assumptions 5, 6 (with r = 1), 7 hold and event A happens, then for all 0 < 32 < 0 the output of EIRM equation 6 Φ̂ = S̃t(α), where α ∈ [ 1
1+ 12ωλmin √ 3 |Etr| 2πmin , 1 1− 12ωλmin √ 3 |Etr| 2πmin ] = [ 1 1+τ √ , 1 1−τ √ ]. Note that
3 2 < 0 =⇒ < th. Hence, all that remains to be shown is that event A occurs with a probability at least 1 − δ. Next, we show that if |D| ≥ 16L ′4
2 log( 2|HΦ| δ ), then with probability 1 − δ event A
happens.
Now we need to bound the probability P(A). We will find an upper bound on the failure probability using Hoeffding’s inequality (Lemma 1) and the bound L′ (derived in equation 77) as follows. We redo the same analysis as was done in equation 54 for reader’s convenience.
P(A) = 1−P ({ D̃ : ∃Φ ∈ HΦ, |R̂ ′ (Φ)−R ′ (Φ)| >
2
}) = P ( ⋃ Φ∈HΦ { D̃ : |R̂ ′ (Φ)−R ′ (Φ)| > 2 })
P ( ⋃
Φ∈HΦ
{ D̃ : |R̂ ′ (Φ)−R ′ (Φ)| >
2
}) ≤ ∑
Φ∈HΦ
P ( {D̃ : |R̂ ′ (Φ)−R ′ (Φ)| >
2
}) ≤ 2|HΦ|e− 2|D| 16L′4
2|HΦ|e− 2|D| 16L′4 ≤ δ =⇒ P(Ac) ≤ δ
|D| ≥ 16L ′4 2 log (2|HΦ|
δ
) =⇒ P(Ac) ≤ δ
(79)
Hence, we know that if |D| ≥ 16L ′4 2 log ( 2|HΦ| δ ) , then with probability 1− δ event A happens.
The proof characterized the property of Φ̂ that satisfies R̂ ′ (Φ̂) ≤ . But how do we know such a Φ̂
exists; we show that a Φ̂ always exists. Consider a Φ that satisfiess the following.
R ′ (Φ) ≤
2
R̂ ′ (Φ)− R̂ ′ (Φ) +R ′ (Φ) ≤
2
R̂ ′ (Φ) ≤
2 + |R̂
′ (Φ)−R ′ (Φ)|
R̂ ′ (Φ) ≤ (follows from event A)
(80)
From equation 73 in the proof of Proposition 11, we know that R ′ (S̃Tγ) = 0. From equation 80, S̃Tγ ∈ HΦ satisfies R̂ ′ (Φ) ≤ .
Since equation 6 is a constrained optimization, a penalty based version (IRMv1) was proposed by Arjovsky et al. (2019) that minimizes R̂(Φ) + λR̂ ′ (Φ). Both Propositions 4 and Proposition 5 can be extended to IRMv1. Below we show the sample complexity analysis for IRMv1 applied to the setting assumed in Proposition 5. Sample complexity analysis of IRMv1 below shows that distance to the OOD solution decays as O( √ 1/λ). Define λth = max{5σ2/3 th, 1}
Corollary 1. For every δ ∈ (0, 1), λ > λth, if Assumptions 5, 6 (with r = 1), 7 hold, and |D| ≥ max { 64L′4λ2 25σ4 log ( 4|HΦ| δ ) , 32L 2λ2 25σ4 log( 4 δ ) }
, then with a probability at least 1 − δ every solution Φ̂ of IRMv1 satisfies Φ̂ = ( S̃Tγ ) α, where α ∈ [ 1
1+τσ √
5 3λ
, 1 1−τσ √
5 3λ
].
Proof. The empirical version of IRMv1 minimizes
R̂(Φ) + λR̂ ′ (Φ)
Let us compute the risk achieved by the ideal invariant predictor S̃Tγ.
R ( S̃Tγ ) = ∑ e∈Etr πeEe [( Y e−γTS̃Xe )2] = ∑ e∈Etr πeEe [( Y e−γTZe1 )2] = ∑ e∈Etr πeEe [ (εe)2 ] = σ2
(81) Define event A: { D : |R̂ ( S̃Tγ ) −R ( S̃Tγ ) | ≤ 2 } . If event A holds, then R̂ ( S̃Tγ ) ≤ σ2 +
2 (82) Define event B: { D̃ : ∀ Φ ∈ HΦ, |R̂ ′ (Φ)−R′(Φ)| ≤ 2 } .
If event B holds, then |R̂′ ( S̃Tγ ) − R′ ( S̃Tγ ) | ≤ 2 and if we plug in R ′(S̃Tγ) = 0 (from equation 73), then
R̂ ′( S̃Tγ ) ≤
2 (83)
Define success asA∩B. If eventA∩B occurs, then from equation 82 and equation 83 the following is true for any solution Φ of IRMv1
R̂(Φ) + λR̂ ′ (Φ) ≤ R̂ ( S̃Tγ ) + λR̂ ′( S̃Tγ ) ≤ σ2 +
2 + λ 2 (84)
Since R̂(Φ) ≥ 0 it follows that
R̂ ′ (Φ) ≤ σ2 + 2 λ + 2 =⇒ R̂ ′ (Φ)−R ′ (Φ) +R ′ (Φ) ≤ σ2 + 2 λ + 2 =⇒ R ′ (Φ) ≤ σ2 + 2 λ +
(85)
In the last implication in the above equation 85, we use the condition that event B occurs. Let = σ 2
λ and substitute in the above equation 85 to get
R ′ (Φ) ≤ 2σ
2
λ +
σ2
2λ2 ≤ 5σ
2
2λ (since λ ≥ 1, σ
2
2λ2 ≤ σ
2
2λ ) (86)
Recall Proposition 11 and see 5σ 2 2λ takes the role of . If 5σ2 2λ ≤ 0 =⇒ λ ≥ 5σ2 3 th , then the condition in the Proposition 11 is true, then every solution of IRMv1 is S̃Tγ(α), where α ∈ [ 1
1+ 1ωλmin √ 5σ2|Etr| 2λπmin , 1 1− 1ωλmin √ 5σ2|Etr| 2λπmin ] = [ 1 1+τσ √ 5 3λ , 1 1−τσ √ 5 3λ ]
We arrived at the above result assuming A ∩ B occurs. We will now show that if |D| ≥ max{ 16L ′4λ2
σ4 log( 4|HΦ| δ ), 8L2λ2 σ4 log( 4 δ )}, then with a probability at least 1− δ, A ∩B occurs. We write P(A) = P ({ D : |R̂ ( S̃Tγ ) − R ( S̃Tγ ) | ≤ 2 }) = 1 − P ({ D : |R̂(S̃Tγ) − R(S̃Tγ)| >
2
}) . From Hoeffding’s inequality in Lemma 1 and 3 it follows that
P ({ D : |R̂(S̃Tγ)−R(S̃Tγ)| >
2
}) ≤ 2 exp ( − |D| 2
8L2
)
2 exp ( − |D| 2
8L2
) ≤ δ
2 =⇒ P(Ac) ≤ δ 2
|D| ≥ 8L 2 2 log (4 δ ) =⇒ P(Ac) ≤ δ 2 (87)
Next, we will show next that if |D| ≥ 16L ′4 2 log( 4|HΦ| δ ), then with probability at least 1− δ 2 event B happens.
Now we need to bound the probability P(B). We will find an upper bound on the failure probability using Hoeffding’s inequality and 3 as follows
P(B) = 1−P ({ D̃ : ∃Φ ∈ HΦ, |R̂ ′ (Φ)−R ′ (Φ)| >
2
}) = P ( ⋃ Φ∈HΦ { D̃ : |R̂ ′ (Φ)−R ′ (Φ)| > 2 } )
P ( ⋃
Φ∈HΦ
{ D̃ : |R̂ ′ (Φ)−R ′ (Φ)| >
2
}) ≤ ∑
Φ∈HΦ
P ({ D̃ : |R̂ ′ (Φ)−R ′ (Φ)| >
2
} ) ≤ 2|HΦ|e− 2|D| 16L′4
2|HΦ|e− 2|D| 16L′4 ≤ δ 2 =⇒ P(Bc) ≤ δ 2 |D| ≥ 16L ′4 2 log (4|HΦ|
δ
) =⇒ P(Bc) ≤ δ
2 (88)
Hence, we know that if |D| ≥ 16L ′4 2 log( 4|HΦ| δ ), then with a probability at least 1 − δ 2 event B happens.
Combining equation 87 and equation 88 we get if |D| ≥ max {16L′4 2 log (4|HΦ|
δ
) , 8L2 2 log (4 δ )} then A ∩B occurs with at least 1− δ probability. Substitute = 5σ 2
2λ and this completes the proof.

heading: 7.4 EXTENSIONS: POLYNOMIAL MODELS, INFINITE HYPOTHESIS CLASSES, BINARY-CLASSIFICATION

heading: 7.4.1 POLYNOMIAL MODEL
text: In the main body of the work, we did not cover the polynomial model in detail due to space limitations. In this section, we study the polynomial model and show how the results for the linear model can be generalized to this case. We first begin by stating the model itself. Assumption 8.
e ∼ Categorical(πe), πe > 0∀e ∈ Etr Y e = γTζcp ( Ze1 ) + εe, εe ⊥ Ze1 ,E[εe] = 0,E[(εe)2] = σ2, |εe| ≤ εsup Xe = S(Ze1 , Z e 2)
(89)
Assume that Ze1 component of S is invertible, i.e. ∃S̃ such that S̃(S(Ze1 , Ze2)) = Ze1 and also S̃tγ 6= 0. In the above ζap is a polynomial feature map of degree p defined as ζap : Ra → Ra ′
, where a denotes the dimension of the input to the map ζap , ζ a p (W ) = [W,W ⊗W, . . . , (W ⊗W...p times⊗ W )] = [(W⊗i)pi=1] and ⊗ is the Kronecker product. Also, a ′ = ∑p i=1 a i. ∀e ∈ Etr, πe ≥ π min
|Etr| . The support of distribution of Ze = [Ze1 , Z e 2 ], PeZe , is bounded and the operator norm of S, ‖S‖ = σmax(S) (σmax(S) is maximum singular value of S), is also bounded.
Define Ze = (Ze1 , Z e 2) and say Z e 1 ∈ Rc and Ze2 ∈ Rd.
Can we directly use the analysis from the linear case? We cannot directly use the polynomial map for the features as we also need to find an appropriate transformation of the matrix S, which preserves the linear relationship between the transformed features and the transformed variables Z. We carry out this exercise below.
Define S̄ = diag [ (S⊗i)pi=1 ] , where S̄ is a block diagonal matrix with diagonal matrices defining it
given as { S, S⊗2 . . . , S⊗p } .
Define X̄e = ζnp ( Xe ) , where ζnp (X
e) is the polynomial feature map of degree p of n dimensional input Xe. Similarly, define Z̄e1 = ζ c p(Z e 1), where ζ c p(Z e 1) is the polynomial feature map of degree p of c dimensional input Ze1 and define Z̄e = ζ c+d p (Z e), where ζc+dp (Z e) is the polynomial feature map of degree p of c+ d dimensional input Ze = (Ze1 , Z e 2). Observe that each component of Z̄e1 is also in Z̄.
From the model we know that Xe = SZe. We would like to remind the reader of the mixed product property of tensors. Consider matrices A ∈ Ri×j , B ∈ Rk×l C ∈ Rj×p, D ∈ Rl×q .
(A⊗B)(C ⊗D) = (AC)⊗ (BD) (90)
In the expressions that follow, we exploit the mixed-product property of Kronecker product stated above.
S̄Z̄e = diag [ (S⊗i)pi=1 ][ (Ze,⊗i)pi=1 ] = [ (SZe)⊗i)pi=1 ] = [ (Xe,⊗i)pi=1 ] = ζnp (X e) = X̄e (91)
Define ¯̃S = diag [ (S̃⊗i)pi=1 ] .
¯̃SXe = ¯̃SS̄Z̄e = diag [ (S̃⊗i)pi=1 ] diag [ (S⊗i)pi=1 ][ (Ze,⊗i)pi=1 ] diag [ ((S̃S)⊗i)pi=1 ][ (Ze,⊗i)pi=1 ] = [ ((S̃SZe)⊗i)pi=1 ] = [ ((Ze1) ⊗i)pi=1 ] = ζcp(Z e 1) = Z̄ e 1
(92)
The dimensionality of X̄e is n ′ = ∑p i=1 n i = n p+1−n n−1 . Assumption 9. Inductve bias. HΦ is a finite set of linear models (bounded) parametrized by Φ ∈ Rn′ . ¯̃STγ ∈ HΦ. ∃ ω > 0,Ω > 0, ∀Φ ∈ HΦ, ω ≤ ‖Φ‖2 ≤ Ω and 2ω ≤ ‖ ¯̃STγ‖2 ≤ 23+2√2Ω,
We next compute the norm of S̄ in terms of the norm of S. Recall we are using operator norm defined as ‖S‖ = σmax(S). ‖S̄‖ = ∥∥∥diag[(S⊗i)pi=1]∥∥∥. Since S̄ is a diagonal matrix ‖S̄‖ = maxi=1,..,p{‖S⊗i‖}. Also, note that‖S⊗i‖ = ‖S‖i (Laub, 2005). Therefore, ‖S̄‖ = maxi=1....,p{‖S‖i}. Hence, if ‖S‖ is bounded, ‖S̄‖ is also bounded.
Also, ‖Z̄e‖2 = ∑ i ‖Ze,⊗i‖2. Observe that ‖Ze,⊗i‖ = ‖Ze‖i. Hence, if ‖Ze‖ is bounded, ‖Z̄e‖ is also bounded. Since X̄e = S̄Z̄e. We can conclude that ‖X̄e‖ is also bounded. We can now follow the same line of reasoning as in equation 74, equation 75,equation 76, equation 77 to conclude that the loss and the gradient of the loss are bounded.
We rewrite the above model in Assumption 8 as a linear model in terms of the transformed features.
e ∼ Categorical(πe), πe > 0∀e ∈ E Y e = γtZ̄1 + ε e, εe ⊥ Z̄1 e ,E[εe] = 0,E[(εe)2] = σ2, |εe| ≤ εsup
X̄e = S̄Z̄e (93)
We showed above in equation 92 that Z̄e1 defined in equation 91 component of Z̄e is invertible, ¯̃SS̄Z̄e = Z̄e1 . We have also shown above that support of Z̄
e is bounded and the norm of S̄ is bounded and as a result ‖X̄e‖ and the loss and the gradient of the loss (conditions in Assumption 3 are satisfied) are bounded. We adapt the linear general position assumption (Assumption 6) for the polynomial case below. Assumption 10. Linear general position of training environments. A set of training environments Etr is said to lie in a linear general position of degree r for some r ∈ N if |Etr| > n
′ − r+n′/r and for non-zero x ∈ Rn ′ and
dim ( span { Ee[X̄eX̄e,T]x− Ee[X̄e,Tεe] } e∈Etr ) > n ′ − r (94)
We denote E[X̄eX̄e,T] = Σ̄e Assumption 11. For all the environments e ∈ Etr, Σ̄e is positive definite.
Define the minimum eigenvalue over all the matrices Σ̄e as λ̄min = mine∈Etr λ(Σ̄e). Define ̄th = 24−16 √ 2
3 πmin |Etr| (ωλ̄min) 2.
From the analysis in this section, we see that we have been able to construct a linear model identical to 5, where the role of Xe, Ze, S, S̃, Σe, λmin, th is taken by X̄e, Z̄e, S̄,
¯̃S, Σ̄e, λ̄min, ̄th. Now we are ready to use the result already proven for linear model and state the next Proposition in terms of the parameters for the polynomial model. Proposition 13. Let ` be the square loss. For every ∈ (0, ̄th), δ ∈ (0, 1), if Assumptions 8, 9, 10 (with r = 1), 11 hold and if the number of data points |D| is greater than 16L ′4 2 log ( 2|HΦ|
δ
) ,
then with a probability at least 1− δ, every solution Φ̂ to EIRM (equation 6) satisfies Φ̂ = ( ¯̃STγ ) α,
where α ∈ [ 1 1+τ √ , 1 1−τ √ ].

heading: 7.4.2 INFINITE HYPOTHESIS CLASSES
text: In the work so far we have assumed that the hypothesis classHΦ is finite. In this section, we discuss infinite hypothesis class extensions. Before we do that we state an important result on covering numbers that we will use soon.
Lemma 5. (Shalev-Shwartz & Ben-David, 2014) Define a set A = {a ∈ Rk, ‖a‖2 ≤ Asup}. Covering number for η-cover of A given as Nη(A) is bounded as
Nη(A) ≤ (2√Asupk
η
)k (95)
A. Infinite Hypothesis Class: Confounders and Anti-causal variables
In this section, we seek to extend Proposition 5 to infinite hypothesis classes.
We restate the Assumption 7 for linear models.
Assumption 12. Inductve bias. HΦ is a set of linear models (bounded) parametrized by Φ ∈ Rn HΦ = {Φ ∈ Rn, 0 < ω ≤ ‖Φ‖2 ≤ Ω}. S̃Tγ ∈ HΦ. 2ω ≤ ‖S̃Tγ‖2 ≤ 23+2√2Ω,
Note that the only difference between Assumption 7 and Assumption 12 is that the hypothesis class is not required to be finite anymore.
We already established in equation 75, equation 76, equation 77 that from Assumptions 5 and 7, we can show that the conditions in Assumption 3 hold, i.e., loss and the gradient of the loss are bounded, and also Xe, Y e are bounded . The same conclusion follows from Assumptions 5 and Assumption 12. Hence, for the rest of this section, we can state that ‖Xe‖ ≤ Xsup, |Y e| ≤ K, the square loss `(Φ(·), ·) is bounded by L and ∂`(w·(Φ(·),·)∂w ∣∣∣ w=1.0 is bounded by L′. In the next lemma, we aim to show that if Assumption 5 and 12 hold, then R ′ (Φ) is Lipschitz continuous.
Lemma 6. If Assumption 5 and 12 hold, then R ′ (Φ) is Lipschitz continuous in Φ.
Proof. The output of the model |ΦTX| ≤ ‖Φ‖‖X‖ ≤ √ ΩXsup (From Cauchy-Schwarz).
R ′ (Φ) = ∑ e πeEe [ ΦTXe ( ΦTXe − Y e )]2 |R ′ (Φ1)−R ′ (Φ2)| =
∣∣∣∑ e πe ( Ee [ ΦT1X e ( ΦT1X e − Y e )]2 − Ee [ ΦT1X e ( ΦT2X e − Y e )]2)∣∣∣
≤ ∑ e πe ∣∣∣Ee[ΦT1Xe(ΦT1Xe − Y e)]2 − Ee[ΦT2Xe(ΦT2Xe − Y e)]2∣∣∣
(96)
We bound each term in the summation in the equation 96 above
∣∣∣Ee[ΦT1Xe(ΦT1Xe − Y e)]2 − Ee[ΦT2Xe(ΦT2Xe − Y e)]2∣∣∣ = ∣∣∣(Ee[ΦT1Xe(ΦT1Xe − Y e)]− Ee[ΦT2Xe(ΦT2Xe − Y e)])(Ee[ΦT1Xe(ΦT1Xe − Y e)]+ Ee [ ΦT2X e ( ΦT2X e − Y e )])∣∣∣
≤ ∣∣∣Ee[ΦT1Xe(ΦT1Xe − Y e)]− Ee[ΦT2Xe(ΦT2Xe − Y e)]∣∣∣2√ΩXsup(√ΩXsup +K)
≤ ∣∣∣Ee[ΦT1Xe(ΦT1Xe − Y e)]− Ee[ΦT2Xe(ΦT1Xe − Y e)]+ Ee[ΦT2Xe(ΦT1Xe − Y e)]− Ee [ ΦT2X e(ΦT2X e − Y e)
]∣∣∣2√ΩXsup(√ΩXsup +K) = ∣∣∣Ee[(ΦT1Xe − ΦT2Xe)(ΦT1Xe − Y e)]+ Ee[ΦT2Xe(ΦT1Xe − ΦT2Xe)]∣∣∣2√ΩXsup(√ΩXsup +K)
≤ ∣∣∣Ee[(ΦT1Xe − ΦT2Xe)(ΦT1Xe − Y e)]∣∣∣2√ΩXsup(√ΩXsup +K)+∣∣∣Ee[ΦT2Xe(ΦT1Xe − ΦT2Xe)]∣∣∣2√ΩXsup(√ΩXsup +K)
(97)
We bound each term in the last line in equation 97∣∣∣Ee[(ΦT1Xe − ΦT2Xe)(ΦT1Xe − Y e)]∣∣∣ ≤ Ee
[∣∣∣(ΦT1Xe − ΦT2Xe)∣∣∣∣∣∣(ΦT1Xe − Y e)∣∣∣] ≤ ( √ ΩXsup +K)Ee [∣∣(Φ1 − Φ2)TXe)∣∣] ≤ ( √ ΩXsup +K)Ee [ ‖Φ1 − Φ2‖Xsup ] (Cauchy-Scwarz) ≤ ( √ ΩXsup +K)Xsup‖Φ1 − Φ2‖
(98)
∣∣∣Ee[ΦT2Xe(ΦT1Xe − ΦT2Xe)]∣∣∣ ≤ Ee
[∣∣∣ΦT2Xe∣∣∣∣∣∣(ΦT1Xe − ΦT2Xe)∣∣∣] ≤ √ ΩXsupEe [∣∣∣(Φ1 − Φ2)TXe)∣∣∣] (Cauchy-Scwarz) ≤ √ ΩXsupEe [ ‖Φ1 − Φ2‖Xsup
] ≤ √ Ω(Xsup)2‖Φ1 − Φ2‖
(99)
Substituting equation 98 and equation 99 in equation 97 to get
|R ′ (Φ1)−R ′ (Φ2)| ≤ 2 √ Ω(Xsup)2( √ ΩXsup +K)(2 √ ΩXsup +K)‖Φ1 − Φ2‖ (100)
Therefore, R ′ is Lipschitz with a constant C ′ ≤ 2 √ Ω(Xsup)2( √ ΩXsup +K)(2 √ ΩXsup +K).
We just showed that R ′ is Lipschitz continuous and we set its Lipschitz constant as C ′ .
Proposition 14. Let ` be the square loss. For every ∈ (0, th) and δ ∈ (0, 1), if Assumptions 5, 6 (with r = 1), 12 hold and if the number of samples |D| is greater than 32L ′4
2
[ n log ( 16C ′√ Ωn ) +
log (
2 δ
)] , then with a probability at least 1 − δ every solution Φ̂ to EIRM (equation 6) satisfies
Φ̂ = S̃Tγ(α), where α ∈ [ 1 1+τ √ , 1 1−τ √ ].
Proof. Following the proof of Proposition 5, our goal is to compute the probability of event A: {∀Φ ∈ HΦ, |R̂ ′ (Φ) − R′(Φ)| ≤ 2}. We construct a minimum cover of size Nη(HΦ) (See Lemma 5) with points C = {Φj}bj=1. Compute the probability of failure at one point Φj in the cover
P [{ D̃ : |R ′ (Φj)− R̂ ′ (Φj)| >
4
}] < 2e− 2|D| 32L′4 (101)
We use union bound to bound the probability of failure over the cover C as follows P [{ D̃ : max
Φj∈C |R
′ (Φj)− R̂ ′ (Φj)| >
4
}] < 2Nη(HΦ)e− 2|D| 32L′4 (102)
Now consider any Φ ∈ HΦ and suppose Φj is nearest point to it in the cover.
|R ′ (Φ)− R̂ ′ (Φ)| = |R ′ (Φ)−R ′ (Φj) +R ′ (Φj)− R̂ ′ (Φj) + R̂ ′ (Φj)− R̂ ′ (Φ)|
≤ |R ′ (Φ)−R ′ (Φj)|+ |R ′ (Φj)− R̂ ′ (Φj)|+ |R̂ ′ (Φj)− R̂ ′ (Φ)| ≤ |R ′ (Φj)− R̂ ′ (Φj)|+ 2ηC ′
(103)
In the above simplification, we exploited the Lipschitz continuity ofR ′ . Therefore, for each Φ ∈ HΦ
|R ′ (Φ)− R̂ ′ (Φ)| ≤ max
Φj∈C |R
′ (Φj)− R̂ ′ (Φj)|+ 2ηC ′
max Φ∈HΦ
|R ′ (Φ)− R̂ ′ (Φ)| ≤ max
Φj∈C |R
′ (Φj)− R̂ ′ (Φj)|+ 2ηC
′ (104)
Set η = 8C′
in equation 104 and from equation 102 with probability at least 1−Nη(HΦ)2e− 2|D| 32L′4
max Φ∈HΦ
|R ′ (Φ)− R̂ ′ (Φ)| ≤ /2 (since max
Φj∈C |R
′ (Φj)− R̂ ′ (Φj)| ≤ /4 ) (105)
We bound Nη(HΦ)2e− 2|D| 32L′4 ≤ δ and solve for bound on |D| to get
|D| ≥ 8L ′4
2
4
log (2Nη(HΦ)
δ
) (Use Lemma 5)
|D| ≥ 32L ′4
2
[ n log (16C ′√Ωn ) + log (2 δ )] (106) Therefore, if condition in equation 106 holds, then event A occurs and following the same argument as in the proof of Proposition 5 the proof is complete.
B. Infinite hypothesis class: Lipschitz continuous functions
In this section, we seek to extend Proposition 2 and 4 to infinite hypothesis class of Lipschitz continuous functions that we formally define next. Define a map Φ : P × X → R from the parameter space P and the feature space X to reals. Each p ∈ P is a possible choice for the representation Φ(p, ·). Consider neural networks as an example, P represents the set of the values the weights of the network can take. Assumption 13. Φ : P × X → R is a a Lipschitz continuous function (with Lipschitz constant say Q).

heading: P ⊂ Rk is closed and bounded, thus there exists a P <∞ such that ∀p ∈ P, ‖p‖2 ≤ P .
text: X ⊂ Rn is closed and bounded, thus there exists a Xsup <∞ such that ∀x ∈ X , ‖x‖ ≤ Xsup. Y ⊂ R is closed and bounded, thus there exists a K <∞ such that ∀y ∈ Y, |y| ≤ K
Assumption 14. Lipschitz loss and gradient of loss. R(Φ) is Lipschitz with a constant C, R ′ (Φ) is Lipchitz with a constant C ′ .
From Assumption 13 derive the conditions in Assumption 3 and Assumption 14 Φ : P×X → R is a continuous function defined over closed and bounded domain P × X (domain is compact) and as a result Φ is bounded say by M .
Consider square loss `(Φ(p,X), Y ) = (Y −Φ(p,X))2 ≤ (M+K)2. Hence, there exists an L such that |`(Φ(·), ·)| ≤ L ≤ (M +K)2. ∂`(w.Φ(·),·)
∂w
∣∣∣ w=1.0 is bounded: ∣∣∣∂`(w.Φ(·),·)∂w ∣∣∣ w=1.0 ∣∣∣ = |(Y − Φ(p,X))Φ(p,X)| ≤ (K + M)M . Hence, there exists an L′ such that |∂`(w.Φ(·),·)∂w ∣∣∣ w=1.0 | ≤ L′ ≤ (K +M)M
Lemma 7. If Assumption 13 holds, then R(Φ(p, ·)) and R′(Φ(p, ·)) are Lipschitz continuous in p.
R is Lipschitz:
|R(Φ(p, ·)−R(Φ(q, ·))| = ∣∣∣∑ e πe ( Ee[(Φ(p,Xe)− Y e)2]− Ee[(Φ(q,Xe)− Y e)2] )∣∣∣ ≤ ∑ e πe ∣∣∣Ee[(Φ(p,Xe)− Y e)2]− Ee[(Φ(q,Xe)− Y e)2]∣∣∣
= ∑ e πe ∣∣∣Ee[Φ(p,Xe)− Φ(q,Xe)]Ee[Φ(p,Xe) + Φ(q,Xe)− 2Y ]∣∣∣
≤ ∑ e πeEe [∣∣∣Φ(p,Xe)− Φ(q,Xe)∣∣∣]Ee[∣∣∣Φ(p,Xe) + Φ(q,Xe)− 2Y ∣∣∣]
≤ ∑ e πeEe [∣∣∣Φ(p,Xe)− Φ(q,Xe)∣∣∣]2(M +K)
≤ ‖p− q‖2(M +K)Q (107)
Therefore, R is Lipschitz with a constant C ≤ 2(M +K)Q
R ′ is Lipschitz:
|R ′ (Φ(p, ·))−R ′ (Φ(q, ·))| = ∣∣∣∑ e πe ( Ee [ Φ(p,Xe) ( Φ(p,Xe)− Y e )]2 − E [ Φ(q,Xe) ( Φ(q,Xe)− Y e )]2)∣∣∣ ≤ ∑ e πe ∣∣∣Ee[Φ(p,Xe)(Φ(p,Xe)− Y e)]2 − E[Φ(q,Xe)(Φ(q,Xe)− Y e)]2∣∣∣
(108)
We bound each term in the summation in the equation 96 above
∣∣∣(Ee[Φ(p,Xe)(Φ(p,Xe)− Y e)]− Ee[Φ(q,Xe)(Φ(q,Xe)− Y e)])(Ee[Φ(p,Xe)(Φ(p,Xe)− Y e)]+ Ee [ Φ(q,Xe) ( Φ(q,Xe)− Y e
)])∣∣∣ ≤ ∣∣∣(Ee[Φ(p,Xe)(Φ(p,Xe)− Y e)]− Ee[Φ(q,Xe)(Φ(q,Xe)− Y e)])∣∣∣2M(M +K)
≤ ∣∣∣(Ee[Φ(p,Xe)(Φ(p,Xe)− Y e)]− Ee[Φ(q,Xe)(Φ(p,Xe)− Y e)]+ Ee[Φ(q,Xe)(Φ(p,Xe)− Y e)]− Ee [ Φ(q,Xe) ( Φ(q,Xe)− Y e
)])∣∣∣2M(M +K) = ∣∣∣(Ee[(Φ(p,Xe)− Φ(q,Xe)(Φ(p,Xe)− Y e)])+ Ee[Φ(q,Xe)(Φ(p,Xe)− Φ(q,Xe))]∣∣∣2M(M +K)
≤ ∣∣∣(Ee[(Φ(p,Xe)− Φ(q,Xe))(Φ(p,Xe)− Y e)])∣∣∣2M(M +K)+∣∣∣Ee[Φ(q,Xe)(Φ(p,Xe)− Φ(q,Xe))]∣∣∣2M(M +K)
(109)
We bound each term in the last line in equation 109∣∣∣(Ee[(Φ(p,Xe)− Φ(q,Xe))(Φ(p,Xe)− Y e)])∣∣∣ ≤ Ee
[∣∣(Φ(p,Xe)− Φ(q,Xe))∣∣∣∣(Φ(p,Xe)− Y e)∣∣] ≤ (M +K)Q‖p− q‖
(110)
∣∣E[Φ(q,Xe)(Φ(p,Xe)− Φ(q,Xe)]∣∣ ≤ E
[∣∣Φ(q,Xe)||(Φ(p,Xe)− Φ(q,Xe)∣∣] ≤MQ‖p− q‖
(111)
Substituting equation 110 and equation 111 in equation 109 to get
|R ′ (Φ(p, ·))−R ′ (Φ(q, ·))| ≤ 2M(M +K)(2M +K)Q‖p− q‖ (112)
Therefore, R ′ is Lipschitz with a constant C ′ ≤ 2M(M +K)(2M +K)Q.

heading: C. EIRM: Sample complexity with no distributional assumptions
text: In this section, we discuss the extension of Proposition 2 to the infinite hypothesis class case. Consider the problem in equation 6 and replace the with + κ. Define distance between the two sets S IV( ) and its approximation S IV( + κ) as follows
dis(κ) = max g∈S IV( +κ) min h∈S IV( ) d(g, h) (113)
where d(g, h) is some metric that measures the distance between functions g and h. Observe that if a ≤ b, dis(a) ≤ dis(b). Assumption 15. limk→0 dis(κ) = 0
Define D∗ = max {32L2
ν2
[ k log (16C√Pk ν ) + log (2 δ )] , 8L′4 κ2 [ k log (8C ′√Pk κ ) + log (2 δ )]} Proposition 15. For every ν > 0 and δ ∈ (0, 1), if Assumption 13, 15 hold, then ∃ κ > 0 such that if the number of samples |D| is greater than D∗, then with a probability at least 1 − δ, every solution Φ̂ to EIRM (replace with + κ inequation 6) in S IV( + 2κ) and |R(Φ̂) − R(Φ∗)| ≤ ν, where Φ∗ is a solution of IRM in equation 5.
Proof. We divide the proof in two parts.
Define an event A: {D̃ : ∀p ∈ P, |R′(Φ(p, ·))− R̂′(Φ(p, ·))| ≤ κ}.
In the first half, we will show that if event A occurs, then S IV( ) ⊆ Ŝ IV( + κ) ⊆ S IV( + 2κ) and then bound the probability of A not occuring.
R ′ (Φ(p, ·)) ≤ =⇒ R ′ (Φ(p, ·))− R̂ ′ (Φ(p, ·)) + R̂ ′ (Φ(p, ·)) ≤ =⇒
R̂ ′ (Φ(p, ·)) ≤ + |R ′ (Φ(p, ·))− R̂ ′ (Φ(p, ·))| =⇒ R̂ ′ (Φ(p, ·)) ≤ + κ
(114)
Therefore, S IV( ) ⊆ Ŝ IV( + κ)
R̂ ′ (Φ(p, ·)) ≤ + κ =⇒ R̂ ′ (Φ(p, ·))−R ′ (Φ(p, ·)) +R ′ (Φ(p, ·)) ≤ + κ =⇒
R ′ (Φ(p, ·)) ≤ + κ+ |R ′ (Φ(p, ·))− R̂ ′ (Φ(p, ·))| =⇒ R ′ (Φ(p, ·)) ≤ + 2κ
(115)
Therefore, Ŝ IV( + κ) ⊆ S IV( + 2κ) We bound the probability of event A not occurring. Using the covering number (from Lemma 5) we construct a minimum cover of size b = Nη(P) with points C1 = {pj}bj=1. Compute the probability of failure at one point pj in the cover
P [{ D̃ : |R ′ (Φ(pj , ·))− R̂ ′ (Φ(pj , ·))| > κ
2
}] < 2e− κ2|D| 8L′4 (116)
We use union bound to bound the probability of the failure over the entire cover C1 as P [{ D̃ : max
pj∈C1 |R
′ (Φ(pj , ·))− R̂ ′ (Φ(pj , ·))| >
κ
2
}] < Nη(P)2e− κ2|D| 8L′4 (117)
Now consider any p ∈ P and suppose pj is nearest point to it in the cover.
|R ′ (Φ(p, ·))− R̂ ′ (Φ(p, ·))| =
|R ′ (Φ(p, ·))−R ′ (Φ(pj , ·)) +R ′ (Φ(pj , ·))− R̂ ′ (Φ(pj , ·)) + R̂ ′ (Φ(pj , ·))− R̂ ′ (Φ(p, ·))|
≤ |R ′ (Φ(pj , ·))− R̂ ′ (Φ(pj , ·))|+ 2ηC ′
(118)
In the above simplficiation, we exploit the Lipschitz continuity of R ′ .
Therefore
∀p ∈ P |R ′ (Φ(p, ·))− R̂ ′ (Φ(p, ·))| ≤ max
pj∈C1 |R
′ (Φ(pj , ·))− R̂ ′ (Φ(pj , ·))|+ 2ηC ′
max p∈P |R
′ (Φ(p, ·))− R̂ ′ (Φ(p, ·))| ≤ max
pj∈C1 |R
′ (Φ(pj , ·))− R̂ ′ (Φ(pj , ·))|+ 2ηC
′ (119)
Set η = κ 4C′
in equation 119 and from equation 117 with probability at least 1−Nη(P)2e− κ2|D| 8L′4
max p∈P |R
′ (Φ(p, ·))− R̂ ′ (Φ(p, ·))| ≤ κ (since max
pj∈C1 |R
′ (Φ(pj , ·))− R̂ ′ (Φ(pj , ·))| ≤ κ/2 ) (120)
We bound Nη(P)2e− κ2|D| 8L′4 ≤ δ and solve for bound on |D| to get
|D| ≥ 8L ′4 κ2 log (2Nη(P) δ ) |D| ≥ 8L ′4
κ2
[ k log (8C ′√Pk κ ) + log (2 δ )] (121)
Therefore, if condition in equation 121 holds, then event A occurs. If event A occurs, then S IV( ) ⊆ Ŝ IV( + κ) ⊆ S IV( + 2κ). Φ(p∗, ·) is a solution to IRM equation 5 and it satisfies ∀Φ(p, ·) ∈ S IV( ) R(Φ(p∗)) ≤ R(Φ(p, ·).
Define an event B: {D : ∀p ∈ P, |R(Φ(p, ·))− R̂(Φ(p, ·))| ≤ ν2}. If event B occurs, then for a solution Φ(p̂, ·) of equation 6 (where is replaced with + κ) satisfies
R(Φ(p̂, ·))− ν 2 ≤ R̂(Φ(p̂, ·) ≤ R̂(Φ(p∗, ·)) ≤ R(Φ(p∗, ·)) + ν
2 R(Φ(p̂, ·)) ≤ R(Φ(p∗, ·)) + ν
(122)
Using the covering number (from Lemma 5) we construct a minimum cover of size b = Nη(P) with points C1 = {pj}bj=1. Let us bound the probability of failure at one point pj in the cover
P [ |R(Φ(pj , ·))− R̂(Φ(pj , ·))| > ν
4
] < 2e− ν2|D| 32L2 (123)
We use union bound to bound the probability defined as P [{ D : max
pj∈C1 |R(Φ(pj , ·))− R̂(Φ(pj , ·))| >
ν
4
}] < Nη(P)2e− ν2|D| 32L2 (124)
Now consider any p ∈ P and suppose pj is nearest point to it in the cover.
|R(Φ(p, ·))− R̂(Φ(p, ·))| = |R(Φ(p, ·))−R(Φ(pj , ·)) +R(Φ(pj , ·))− R̂(Φ(pj , ·)) + R̂(Φ(pj , ·))− R̂(Φ(p, ·))| ≤ |R(Φj)− R̂(Φj)|+ 2ηC
(125)
Therefore, for each p ∈ P
|R(Φ(p, ·))− R̂(Φ(pj , ·))| ≤ max pj∈C1 |R(Φ(pj , ·))− R̂(Φ(pj , ·))|+ 2ηC
max p∈P |R(Φ(p, ·))− R̂(Φ(p, ·))| ≤ max pj∈C1 |R(Φ(pj , ·))− R̂(Φ(pj , ·))|+ 2ηC
(126)
Set η = ν8C in equation 124 and from equation 126 with probability at least 1−Nη(P)2e − ν 2|D| 32L′4
max p∈P |R(Φ(p, ·))− R̂
′ (Φ(p, ·))| ≤ ν/2 (since max
pj∈C1 |R
′ (Φ(pj , ·))− R̂ ′ (Φ(pj , ·))| ≤ ν/4 ) (127)
We bound Nη(P)2e− ν2|D| 32L2 ≤ δ and solve for bound on |D| to get
|D| ≥ 32L 2 ν2 log (2Nη(Φ) δ ) |D| ≥ 32L 2
ν2
[ k log (16C√Pk ν ) + log( 2 δ ) ] (128)
Therefore, if condition in equation 128 holds, then with probability at least 1− δ2 event A occurs. Also, if
|D| ≥ max {32L2
ν2
[ k log (16C√Pk ν ) +log (2 δ )] , 8L′4 κ2 [ k log (8C ′√Pk κ ) +log (2 δ )]} (129)
then the event A ∩B occurs with at least 1− δ probability.
Project Φ(p̂, ·) ∈ Ŝ IV( + κ) on S IV( ), i.e., find the closest function in terms of the metric dis, to obtain Φ(˜̂p, ·). If event A occurs, then p̂ ∈ S IV( + 2κ). The distance ‖p̂− ˜̂p‖ ≤ dis(2κ).
|R(Φ(p̂, ·))−R(Φ(˜̂p, ·))| ≤ C‖p̂− ˜̂p‖ ≤ Cdis(2κ) (130)
We choose κ0 such that κ < κ0 (use Assumption 15) such that Cdis(2κ) ≤ ν.
R(Φ(p∗, ·) ≤ R(Φ(˜̂p, ·)) ≤ R(Φ(p̂, ·)) + ν (131)
Therefore, by combining equation 122 amd equation 131, we can conclude that if event A ∩ B occurs, then Φ(p̂, ·) ∈ S IV( + 2κ) and |R(Φ(p∗, ·)) − R(Φ(p̂, ·))| ≤ ν. From the conditions on |D|, we know that A ∩B occurs with probability 1− δ. We substitute Φ(p∗, ·) as Φ∗ and Φ(p̂, ·) as Φ̂ and this completes the proof.

heading: D. OOD Performance: Covariate shift case
text: In this section, we discuss the extension of Proposition 4 to the infinite hypothesis class case. Define D∗1 = { 32L2
ν2
[ k log ( 32C √ Pk
ν
) + log ( 4 δ )] , 16L ′4 2 log ( 4 δ )} .
Define D∗2 = 32L2
ν2
[ k log ( 16C √ Pk
ν
) + log ( 2 δ )] .
Proposition 16. If Assumptions 4, 13 hold, m ∈ HΦ and if the number of samples |D| is greater than D∗1 , then with a probability at least 1− δ, every solution Φ̂ to EIRM satisfies R(m) ≤ R(Φ̂) ≤ R(m) + ν.
If Assumptions 4, 13 hold, m ∈ HΦ and if the number of samples |D| is greater than D∗2 , then with a probability at least 1− δ every solution Φ† of ERM satisfies R(m) ≤ R(Φ†) ≤ R(m) + ν.
Proof. We begin with the first part. Following the proof of Proposition 5, our goal is to compute the probability of event A: {∀p ∈ P, |R̂(Φ(p, ·)) − R(Φ(p, ·))| ≤ ν2}. Using the covering number (from Lemma 5) we construct a minimum cover of size b = Nη(P) with points C = {pj}bj=1. Compute the probability of failure at one point pj in the cover
P [{ D : |R(Φ(pj , ·))− R̂(Φ(pj , ·))| > ν
4
}] < 2e− ν2|D| 32L2 (132)
We use union bound to bound the probability of failure over the cover C
P [{ D : max
j |R(Φ(pj , ·))− R̂(Φ(pj , ·))| >
ν
4
}] < Nη(P)2e− ν2|D| 32L2 (133)
Now consider any p ∈ P and suppose pj is nearest point to it in the cover.
|R(Φ(p, ·))− R̂(Φ(p, ·))| = |R(Φ(p, ·))−R(Φ(pj , ·)) +R(Φ(pj , ·))− R̂(Φ(pj , ·)) + R̂(Φ(pj , ·))− R̂(Φ(p, ·))| ≤ |R(Φ(pj , ·)− R̂(Φ(pj , ·))|+ 2ηC
(134)
In the above simplification, we used the Lipschitz continuity of R. Therefore
∀p ∈ P, |R(Φ(p, ·))− R̂(Φ(p, ·))| ≤ max pj∈C |R(Φ(pj , ·))− R̂(Φ(pj , ·))|+ 2ηC
max p∈P |R(Φ(p, ·))− R̂(Φ(p, ·))| ≤ max pj∈C |R(Φ(pj , ·))− R̂(Φ(pj , ·))|+ 2ηC
(135)
Set η = ν8C in equation 135 and from equation 133 with probability at least 1−Nη(P)2e − ν2n 32L′4
max p∈P |R
′ (Φ(p, ·))− R̂ ′ (Φ(p, ·))| ≤ ν/2 (since max
pj∈C |R
′ (Φ(pj , ·))− R̂ ′ (Φ(pj , ·))| ≤ ν/4 ) (136)
We bound Nη(P)2e− ν2|D| 32L2 ≤ δ and solve for bound on |D| to get
|D| ≥ 8L 2
ν2
4
log (Nη(P)
δ
) (Use Lemma 5)
|D| ≥ 32L 2
ν2
[ k log (16C√Pk ν ) + log (2 δ )] (137) Therefore, if condition in equation 137 holds, then eventA occurs. Observe that the optimal solution p∗ for expected risk minimization p∗ ∈ arg minp∈P R(Φ(p, ·)) satisfies Φ(p∗, ·) = m(·) (From Lemma 4 and m ∈ HΦ). If event A occurs, then from same argument used in equation 48, a solution p+ ∈ P of ERM satisfies R(Φ(p∗, ·)) ≤ R(Φ(p+, ·)) ≤ R(Φ(p∗, ·)) + ν is true. We now move to the second part. From the first part of the proof, we conclude that when
|D| ≥ 32L 2
ν2
[ k log (16C√Pk ν ) + log( 4 δ ) ]
(138)
with a probability at least 1− δ2 event A occurs.
Define an event B: D̃ is such that |R̂′(m)− R′(m)| ≤ 2 . Since R ′ (m) = 0, |R̂′(m)− R′(m)| ≤ 2 =⇒ |R̂ ′ (m)| ≤ 2 =⇒ R̂ ′ (m) ≤ 2 . Therefore, m ∈ Ŝ IV( ).
We write P(B) = P({D̃ : |R̂′(m) − R′(m)| ≤ 2}) = 1 − P({D̃ : |R̂ ′ (m) − R′(m)| > 2}). The gradient of loss function is bounded |∂`(Φ(·),·)∂w |w=1.0| ≤ L ′ . From Hoeffding’s inequality in Lemma 1 it follows that
P(|D̃ : R̂ ′ (h)−R ′ (h)| >
2 ) ≤ 2 exp
( − |D| 2
16L′4
)
2 exp ( − |D| 2
16L′4
) ≤ δ
2
|D| ≥ 16L ′4 2 log (4 δ ) (139)
Combining the two conditions equation 138 and equation 139,
|D| ≥ max {32L2
ν2
[ k log (32C√Pk ν ) + log (4 δ )] , 16L′4 2 log (4 δ )} This ensures P (A ∩B) ≥ 1− δ. If event A ∩B occurs, then we follow the same justification as in the proof of Proposition 4 to claim that a solution p̂ ∈ P to EIRM equation 6 satisfiesR(Φ(p∗, ·)) ≤ R(Φ(p̂, ·)) ≤ R(Φ(p∗, ·)) + ν This completes the proof.

heading: 7.4.3 EXTENSIONS TO BINARY CLASSIFICATION (CROSS-ENTROPY)
text: In the main body of the manuscript, we focused on regression (square-loss). In this section, we discuss the results that can be extended to binary classification (cross-entropy) loss. We will not go in the order in which the results were introduced in the manuscript but in an order that makes for easier exposition for the classification case.
We begin by showing how to extend Proposition 4 to binary classification (cross-entropy). Recall that the entropy of a distribution PX is H(P) = −EP [ log(dP) ] . Recall that the cross entropy of
Q relative to P is H(P,Q) = −EP [ log(dQ) ]
= H(P ) + KL(dP‖dQ). The cross entropy loss ` for binary classification when using a predictor f : X → [0, 1] (f(Xe) is the probability of label 1 conditional on Xe) is given as ` ( f(Xe), Y e ) = Y e log ( f(Xe) ) + (1 − Y e) log ( 1 − f(Xe) ) . For the discussion below Q(Y e|Xe) is defined in terms of f as follows Q(Y e = 1|Xe) = f(Xe)( Q(Y e = 0|Xe) = 1− f(Xe) ) .
Re(f) = Ee [ `(Y e, f ( Xe ) ) ]
= Ee [ Y e log ( f(Xe )) + (1− Y e) log ( 1− f ( Xe ))]
= Ee [ Ee [ Y e|Xe ] log ( f ( Xe )) + ( 1− E [ Y e|Xe ]) log ( 1− f ( Xe ))]
= Ee [ P(Y e = 1|Xe) log ( f ( Xe )) + P(Y e = 0|Xe) log ( 1− f ( Xe ))]
= Ee [ H ( P(Y e|Xe),Q(Y e|Xe) )] = Ee [ H ( P(Y e|Xe) ) + KL ( P(Y e|Xe)‖Q(Y e|Xe
)] = Ee [ H ( P(Y e|Xe) )] + Ee [ KL ( P(Y e|Xe)‖Q(Y e|Xe )]
(140)
From the above it is clear that Q(Y e|Xe ) = P(Y e|Xe )
minimizes the risk in an individual environment. Assumption 16. Invariance w.r.t all the features. For all e, o ∈ Eall and for all x ∈ X , E[Y e|Xe = x] = E[Y o|Xo = x]. Xe ∼ PeXe and ∀e ∈ Eall support of PeXe is equal to X .
Observe that in the binary-classification setting the above assumption amounts to equating the conditional probabilities P(Y e|Xe) and P(Y o|Xo). Recall that map m (from equation 2) simplifies to ∀x ∈ X
m(x) = Ee[Y e|Xe = x] = P(Y e = 1|Xe = x) (141)
If Assumption 16 holds, then from cross-entropy decomposition in equation 140 it is clear that m solves the OOD problem (as it is optimal w.r.t each environment). It is also the unique minimizer. We can justify it based on the same argument presented in equation 56. Suppose there was another optimizer which was different from m over a set with a non-zero measure. Over such a set the the KL divergence term inside equation 140 will be greater than zero, thus making the second term in equation 140 positive thus contradicting the optimality. This shows m is the unique optimizer. The rest of the arguments presented in the proof of Proposition 4 carry over to this case. Therefore, Proposition 4 extends to the cross-entropy loss.
Note that Proposition 2’s proof was agnostic to loss type and only used boundedness, which holds for both cross-entropy as long as the probability output are in the strict interior of [0, 1] defined by [pmin, pmax] ⊂ [0, 1]. We could not generalize Proposition 5 to cross-entropy loss and that is left as future work.
Next, we move to showing how Proposition 1 can be generalized to binary classification. Assumption 17. Existence of an invariant representation. ∃ Φ∗ : X → Z such that ∀e, o ∈ Eall and ∀x ∈ X , E[Y e|Φ∗(x)] = E[Y o|Φ∗(x)].
Recall m defined in equation 2, ∀z ∈ Φ∗(X ) m(z) = Ee[Y e|Ze = z] = P(Y e = 1|Xe = z) (142)
Define a composite predictor w ◦Φ∗. Substituting f = w ◦Φ∗ in equation 140 we get the following. For the discussion below, a distribution R(Y e|Xe) is defined in terms of w ◦Φ∗ as follows R(Y e = 1|Xe) = w ◦ Φ∗(Xe) (R(Y e = 0|Xe) = 1− w ◦ Φ∗(Xe)).
Re(w ◦ Φ∗) = Ee [ H ( P(Y e|Ze) )] + Ee [ KL ( P(Y e|Ze)‖R(Y e| ( Ze ))]
(143)
If all the data is transformed by Φ∗, then from the above decomposition equation 143 it is clear that R(Y e|Ze) = P(Y e|Ze) is the optimal predictor for each environment. Hence, w∗(Ze) = P(Y e = 1|Ze) is the best choice for w. Assumption 18. Existence of an environment where the invariant representation is sufficient. ∃ an environment e ∈ Eall such that Y e ⊥ Xe|Ze, where Ze = Φ∗(Xe).
We derive a relationship as follows for the environment q satisfying Assumption 18.
P(Y q|Zq, Xq) = P(Y q|Zq) (follows from conditional independence in Assumption 18) (144)
Also, note that since Zq = Φ∗(Xq) we have
P(Y q|Zq, Xq) = P(Y q|Xq) (145)
From equation 144 and equation 145 we have
P(Y q|Xq) = P(Y q|Zq) (146)
We use equation 146 in the cross entropy decomposition from equation 140 Rq(f) = Eq [ H ( P(Y q|Zq) )] + Eq [ KL ( P(Y q|Zq)‖Q(Y q|Xq )] (147)
Recall Q(Y q = 1|Xq) = f(Xq) (Q(Y q = 0|Xq) = 1 − f(Xq)). Also, recall w∗(Zq) = P(Y q = 1|Zq). From the above it is clear that f = w∗ ◦ Φ∗ is the optimal predictor for environment q.
Rq(w∗ ◦ Φ∗) = Eq [ H ( P(Y q|Xq) )] (148)
The expected conditional entropy for environment e is defined as H̄e = Ee [ H ( P(Y e|Ze) )] is the
risk achieved by w∗ ◦Φ∗. Also, H̄e measures the amount of noise in the environment. This is much like the variance that remains in the least squares minimization. In the next assumption, we state that the noise in all the environments is bounded above. We also assume that one of the environments which achieves the maximum noise level is environment q, which satisfies Assumption 18. Assumption 19. ∀e ∈ Eall, H̄e ≤ H̄sup, H̄q = H̄sup
Therefore, Rq(w∗ ◦ Φ∗) = H̄sup (149)
From equation 143 for all the environments
Re(w∗ ◦ Φ∗) = H̄e (150)
Observe that maxe∈Eall R e(w∗ ◦ Φ∗) = H̄sup. From the above assumption it is clear that for all predictors f : X → [0, 1].
∀f, max e∈Eall Re(f) ≥ Rq(f) ≥ H̄sup
min f max e∈Eall
Re(f) ≥ H̄sup (151)
Since maxe∈Eall R e(w∗ ◦ Φ∗) = H̄sup, we conclude that w∗ ◦ Φ∗ is the predictor that solves the OOD problem in equation 1. This completes the extension of Proposition 1 to cross-entropy.

heading: 7.4.4 ON THE BIASEDNESS OF ERM
text: Consider the model in Assumption 5. For each environment e ∈ Etr, define a vector ρe = Ee [ εeXe ] . Define a matrix ρ̄ with ρe as column vectors ρ̄ = [ρ1, . . . , ρ|Etr|]. Define a vector π̄ = [π1, . . . , π|Etr|], where recall from Assumption 5 πo is probability a point comes from environment o. Proposition 17. If Assumption 5 holds,HΦ is a linear hypothesis class with parameter Φ and if the rank of ρ̄ is at least one, then ERM is asymptotically biased, i.e., even with infinite data ERM will not achieve the desired solution S̃Tγ, except over a set of measure zero of probability distributions π̄.
Proof. Consider the case when ERM has access to infinite data, i.e., we are solving the expected risk minimization problem stated as minΦ∈HΦ R(Φ). We will consider the linear model in Assumption 5 and assume HΦ linear hypothesis class parametrized by Φ ∈ Rn. We simplify the ∇ΦR(Φ) for the square loss below
∇ΦR(Φ) = ∑ e∈Etr πeEe [( Y e − ΦTXe ) Xe ]
(152)
We compute the gradient for Φ = S̃Tγ as ∇Φ|Φ=S̃TγR(Φ) = ∑ e∈Etr πeEe [( Y e − γTS̃Xe ) Xe ] (Use Assumption 5)
= ∑ e∈Etr πeEe [ εeXe ] (153) Recall ρe = Ee [ εeXe ] and ρ̄ = [ρ1, . . . , ρ|Etr|]. Recall π̄ = [π1, . . . , π|Etr|]. Setting the gradient defined in equation 153 to zero and using the above matrix notation we get
ρ̄π̄ = 0,1Tπ̄ = 1, π̄ ≥ 0 (154)
If π̄ satisfies equation 154, then ERM is unbiased, else it is not. Consider the set of vectors in the probability simplex {π̄ | 1Tπ̄ = 1, π̄ ≥ 0} and define a uniform probability distribution over it. Since rank of ρ̄ > 0 at least one of the columns of ρ̄ is non-zero. As a result a uniform random draw from this set of probablity distributions would have zero probability of satisfying ρ̄π̄ = 0. Therefore, S̃Tγ is not the optimal solution to ERM and thus the solution of ERM would be biased away from S̃Tγ.
In Proposition 5, we had assumed Assumptions 5, 6, 7 hold. If we also assume that rank of ρ̄ is at least one, the Proposition 5 continues to hold. If for at least one e ∈ Etr, Ee [ εeXe ] is non-zero, then the rank of ρ̄ is at least one. From proof of Theorem 10 in Arjovsky et al. (2019), linear general position continues to hold except over a set of covariance matrices with measure zero even when one of the Ee [ εeXe ] is non-zero.
Also, in the above Proposition 17, we only required that rank of ρ̄ is at least one. However, if we make the additional assumptions 5, 6, 7, the result of the above Proposition continues to hold. Therefore, if Assumption 5, 6, 7 hold and rank of ρ̄ is at least one, then ERM is asymptotically biased and IRM can be within √ neighborhood of the ideal solution with the sample complexity shown in Proposition 5.

Review summaries:
The paper investigates the choice of learning paradigms to reach out-of-distribution generalization, namely IRM vs ERM under different scenarios of domain generalization. Technically, generalization bounds and rates are calculated to be able to compare theoretically how each paradigm fares in the different scenarios. Proposed analysis shows that IRM generalizes better than ERM in a number of important cases: presence of confounders and/or anti--causal variables. Numerical experiments on variants of the Colored MNIST benchmark (for each scenario) verify the theoretical findings.

This paper considers a learning scenario where training data $(X,Y)$ comes from a mixture, where membership in each mixture component (“environment”) is clearly labeled. Generalization is required not just under the same mixture, but potentially under changing mixing distributions. This is captured via several alternative formulations. The key underlying assumption that makes the extrapolation possible is the existence of a a representation $\Phi(X)$ that leads to an environment-invariant statistics (by focusing on quadratic loss and linear predictors, the paper asks for invariant $E[Y|\Phi(X)]$ and constant $var[Y|\Phi(X)]$). A range of other assumptions are made to enable analysis, such as sufficiency of the representation in at least one environment and boundedness of the loss and its gradient.