Empirical or Invariant Risk Minimization? A Sample Complexity Perspective
Authors: Kartik Ahuja
Jun Wang
Amit Dhurandhar
Karthikeyan Shanmugam
Kush R. Varshney
Venue: ICLR 2021 Poster (2021)
Keywords: invariant risk minimization
IRM

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

Published as a conference paper at ICLR 2021
EMPIRICAL OR INVARIANT RISK MINIMIZATION? A SAMPLE COMPLEXITY PERSPECTIVE
Kartik Ahuja', Jun Wang", Amit Dhurandhar', Karthikeyan Shanmugam', Kush R. Varshney! f IBM Research, TJ Watson Research Center, NY, Rensselaer Polytechnique Institute
ABSTRACT
Recently, invariant risk minimization (IRM) was proposed as promising solu- tion to address out-of-distribution (OOD) generalization. However, it is unclear when IRM should be preferred over the widely-employed empirical risk mini- mization (ERM) framework. In this work, analyze both these frameworks from the perspective of sample complexity, thus taking a firm step towards an- swering this important question. We find that depending eration 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 ap- proaches arrive at different asymptotic solutions where IRM is guaranteed be close to the desired OOD solutions in the finite sample regime for polynomial gen- erative models, while ERM 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.
the type of data gen-
INTRODUCTION
A recent study shows that models trained to detect COVID-19 from chest radiographs rely spuri- ous 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 hold outside specific train- ing distribution. In one commonly cited example, |Beery et al. (2018) trained 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. (0199 proposed framework called invariant risk minimization (IRM) to address the problem 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 many problems requiring OOD generalization (Gulrajani & Lopez- Paz| (202)). 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) * ptest (X,Y)) that we consider informally stated satisfy an in-
variance condition - there exists a representation &- of the covariates such that patrain (Y (xX("(()) = prest 45
A special case of this occurs when th is identity - Iptrain (X) o ( (X)) = P" prest (Yl such shift is known as covariate-shif (Gretton et al.
(test(X) but (Y
2009). In many other settings Ø* may not be identity (denoted as I), examples include settings with confounders anti-causal variables (Pearl 2009) where covariates appear spuriously corre- lated with the label and pptrain [y * ptesr (Yl)) We use causal Bayesian networks to illustrate these shifts in Figure [ Suppose X [ represents the image, where Xi is the shape of the animal and XG is the background color, Y is the label of the animal, and e is the index of the environment/domain. In Figure [p) X is independent of (Ye, Xf), it represents the covariate shift

---

Published as a conference paper at ICLR 2021
D). In Figure [p) X is spuriously correlated with y through the confounder g², In X is spuriously correlated with Y as it is anti-causally related to Y In both Figure [pp and c) d # I; t0 is block diagonal matrix that selects Xf.
case ($+
Figure T
Our setup assumes we are given data from multiple training environments satisfying the invariance condition, i.e., ((Y|(=(X),
the same across all of them. Ideally, we want to learn and predict show later where prove
using || + (X)]; this predictor has a desirable OOD behavior min-max optimality with respect to (w.r.t)) unseen test distributions satisfying the invariance condi- tion. Our goal is to analyze and compare ERM and IRM's ability learn E||| (X)) from finite training data acquired from a fixed number of training environments. Our analysis has two parts. 1) Covariate shift case ($*
D: ERM and IRM achieve the same asymptotic solution B] We prove (Proposi- 4ion[4) that the sample complex- ity for both the methods is sim- ilar thus there no clear win- ner between the two in the finite sample regime, For the setup in Figure Ta). both ERM and IRM learn a model that only uses Xf.
es
X
Ee
X
g
X
X
X
y
X
Environment for covariate shift: No clear vinner (Prapos toon 4)
b) Environment for confounder- based shitt: IRM wins
Ervironment ans causal shift: (Propositon S and
(Proposition 5 and 6)
Figure Causal Bayesian networks for different distribution shifts.
2) Confounder/Anti-cauuse
variable case (Φ* # I): We consider 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| ( (X)]- We prove that IRM can learn a solution that is within O( distance from E| ( (X)] with sample complexity that increases as O ((4) 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 [T) and c), IRM gets close only using Xf, while ERM even with infinite data (Proposition 17 in the supplement) continues to use X5. We summarize the results in Table !
Arjovsky et al. (2119) proposed the colored MNIST (CMNIST) dataset; comparisons on it showed how HRMeb 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 ad- dition to the original one that map to different real-world settings: i) covariate shift based CMNIST (CS-CMNIST): relies selection bias to induce spurious correlations, ii) confounded CMNIST (CF-CMNIST): relies on confounders to induce spurious correlations, iii) anti-causal CMNIST (AC- CMNIST): this is the original CMNIST proposed by Arrovsky et al. (011)), 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 th * class de- scribed above, IRM has 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.
2 RELATED WORKS
IRM based works. Following the original work IRM from Arjovsky et al. (2019) there have been several interesting works (Teney al., 2020 Krueger et al. [2000] Ahuja et al. 2020 Chang ctt 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 al. (019) 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 (Ben- David 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 Mat- suura & Harada 2020: Deng et al. 2020 David et al. 2010 Pagnoni et al. 2018) is an incomplete representative list (see Redko al. (2019) for further references)
that build the theory of do- main adaptation and generalization and construct new methods based on it. While many of these
works develop bounds loss over the target domain using train data and unlabeled target data,
2

---

Published as a conference paper at ICLR 2021
Table 1: Summary of (empirical) IRM vs. ERM for finite hypothesis class Hg. e: slack in IRM con- straints, V: approximation w.r.t optimal risk, 8: failure probability, Etr: set of training environments, n: data dimension, p: degree of the generative polynomial, L, L': bound on loss & its gradients.
Assumptions
Method
Sample complexity
OOD
ERM
SL- log
2 He
Yes
Covariate shift case: [X* is invariant (Proposition
4
IRM
max { 1822 log 1 (1HHll) ), - log (2)
Yes
Confounder/Anti-cuusa variable case: E E (X=])(((X*)) is invariant, Linear, Polynomial models,
ERM
82 log (² (HHl)
No
IRM
161 log Hol
Yes
|Erll = Onn)) (Proposition
6 17)
some (Ben-David & Urner 2012, David et al., 2010 Pagnoni al. 2018) analyze the finite sample (PAC) guarantees for domain adaptation under covariate shifts. These works (Ben-David & Umnr 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 c.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.
3
SAMPLE COMPLEXITY OF INVARIANT RISK MINIMIZATION
3.1 INVARIANT RISK MINIMIZATION
We start with some background on IRM AArovssy et al. 2019). Consider dataset D = { Jeef... which is a collection of datasets D ments Ex,, where is the index of the environment, i is the index of the data point in the en- vironment, ne is the number of points from environment, of E X C Rn is the feature value and yf E V S R is the corresponding label. Define probability distribution fx" Jeee.. x² is the probability that a training data point is from environment e. Define a probability distribution of points conditional on environment e as IP, (X*, Y*) Pe. Define the joint distribution IP, (Xe,ye,e) ~ P. dip(Xe, Ye,e)= n dPe (Xe, Y²). D is a collection of i.i.d. samples from IP. De- fine 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 (f) = E [xx)(x), Y"), where is the loss, f(X*) is the predicted value, Y the corresponding label and E is the expectation conditional on environment e. The overall expected risk across the training environments is R(f) Eese.r T" R (f). We are interested in two settings: regression (square loss) and binary- sinary-llasiffiattion (cross-entropy loss). In the main body, our focus is regression (square loss) and we mention wherever the results extend to binary-llassfficttion (cross-entropy). We discuss these extensions in the supplement.
{0 3/ ==
obtained from set of training environ-
OOD generalization problem. We want to construct a predictor f that performs well across many unseen environments Eatl. where Eall 2 Ey, For o E Eaul E,., the distribution Ip can be very different from the train environments. Next we state the OOD problem.
min max Rd (f) JEF eEfall
(I)
The above problem is very challenging to solve since we only have access to data from training environments &rr but are required to find the robust solution over all environments Eatl. Next, we make assumptions Eaul and characterize the optimal solution to equation
Assumption Invariance condition. There exists a representation 4- that transforms X to Z = Φ (Xx)) and Ve, o E Eall. Vz E (x') satisfies E"[Y"| IZ = z] E ZY = 2). Also, Ve E Eait, Vz $* (2)', Var [Y = =², where Var is the conditional voriance.
The above assumption is inspired from causality (Pearl 2009). 4- acts as the causal feature extractor and from the definition of causal features, it follows that E [ 2] does not vary across
3

---

Published as a conference paper at ICLR 2021
environments. When a human labels 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 Ye, 0 E Eall. Va E X, P(Y+Ixe = x) and PPY x = x)) are equal thus implying zi. Therefore, for covariate shifts, g5 is identity in Assumption 1
[+ | x]]
A simple instance illustrating Assumption D with o = is when Y < ((X")+"" where [E"
0, (E(((")) = o², E" 1 X Using Assumption
we define the invariant map m : (x(x) IR as
follows
Vz E (X)), m(2) E [ = 2],
where Z +x*
(2)
Assumption 2. Existence of an environment where the invariant representation is sufficient. an environment e Eill such that Y 1 2
Assumption states there exists an environment where the information that X has about ye is also
contained in Z Define composition कh, Vz E X, (x) E. IYee o0(x))
Proposition 1. If is the square loss, and Assumptions and 2 hold, then m o th solves the 0OD problem (equation [
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, focus learning m th as it solves the OOD problem. For covariate shifts th = m (z) = E. [Y = x]] is the OOD solution. In Arjovsky et al. (2019), proof connecting m x and OOD was not stated. Recently, in Koyama & Yamaguchi (020),, result similar to Proposition D was shown but with few differences. The authors assume conditional probabilities are invariant unlike our assumption that only requires con- ditional expectations and variances to be invariant. However, their result applies to more losses. m o 4- is the target we want to learn. Arjovsky et al. (2019) proposed IRM since standard min-max optimization over the training environments E,r and ERM fail to learn m 6" in many cases. The authors in (Arovssy et al. (2019) identify a erucial property of m t and use it to define an object called invariant predictor that we define next.
Invariant predictor and IRM optimization. Define a representation map 4 : X + from feature space to representation space Z C RR, Define classifier map, Z + IR from representation space to real values. Define Hs and Hw as the spaces of representations and classifiers respectively. A data representation & elicits an invariant predictor w ① across environments is a classifier to that achieves the minimum risk simultaneously for all the environments, i.e., Ve E Exr, w E arg min R (w o b) Observe that if we we transform the data with representation 4' then m will achieve the minimum risk simultancously in all the environments. If q- e H4 and m Har then m d- 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:
Ee, if there
min WE
R(w D)
Σ
&' R'
BEE
(3)
s.t. W E arg min Rd (w o D), Ve E Esr
We Hi
From the above discussion we know m 6' is feasible solution to equation It is also the ideal
solution we want IRM find since it solves equation [ Later in Propositions and 6 we show
that IRM actually solves equation equation [ For the setups in Proposition and 6 conventional ERM based approaches fail thus justifying the need for above formulation.
3.2 SAMPLE COMPLEXITY OF GRADIENT CONSTRAINT FORMULATION OF IRM
In Arjovsky et al. (2019). a gradient constrained alternate (derived below in equation 4 to equation was proposed, which focuses on linear and scalar classifiers (Z = R. X - IR. H. RR. In this case, the composite predictor w t is multiplication of w and Tb written as w T. (For binary- classification predictor's output w- no-O (xx) represents logits.) From the definition of invariant predictors and H., = RI it follows that if Vă E I. R (1 4) R (Ø 4). then & is invariant predictor. For square and cross-entropy losses, (' ②) is convex in w. Therefore, a gradient constraint [/ Rd (w- [I = 0 is equivalent to the condition that Vw R, (1.1 (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 gradient constrained optimization as follows
3
4

---

Published as a conference paper at ICLR 2021
min R(Φ) +E
(4)
s.t. Vulw=LoR [.] = 0, Ve &tr
We propose an approximation of the ahove with slack in the constraint. Define R (Φ)
Eeee, V www==lufff (w- and set (e) { R' (+) S & Ha). Note that R is
very similar to the penalty defined in Arjovsky et al. (2019). The approximation of equation is
min
R(
(5)
(ES
c 0, then equation 4 and equation are equivalent. In all the optimizations so far, the pectations are computed w.r.t. the distributions P which are unknown. Therefore, we develop an empirical version of qquationSbbllov (in equation and call it empirical IRM (EIRM). We replace R and R with empirical estimators R and R respectively. For R we use 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 set SIV (e) = { R (x) S e, & Hz). min R((Φ) (c))
we replaced SIV (e) with He in equation 6 then we get the standard ERM. ERM aims to solve min 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. 3 < xc, < 00 such that
VI E Ho. Vx E N, E y/ (z)) v) S T., (f(w--((rr..)) Ow
1w=1ol S T.'.
If every + in the hypothesis class Ho is bounded by M and the label space Y is bounded, then for both square and cross-entropy loss, E()((),.. and ot(ar-(()..)) dhe w==10 are bounded. Define & min R (+) c|: measures how close any penalty can get to the boundary K quantifies how good the finite sample approximation R need to be in order to get (c) (c)()). Define to quantify the approximation w.r.t. optimal risk.
Proposition 2. For every > 0, > 0 and 5 E (0, 1), if is finite hypothesis class, Assumption holds, k > O, and if the number of samples |D] is greater than max 16L SLa
log ( (Hllfe))
then with probability at least 1 - 6, every solution & to EIRM (equation 6 is approximation of IRM, i.e. & E (e), R(**)= - R(()= < R(4*) where 4 is solution to IRM (equation 5
Proof Sketch. The standard analysis in learning theory on ERM or regularized/cconttrined ERM typically relies on linearly separable loss functions. In such cases, we can use standard plug- in 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( + V we first need to guarantee that the set of invariant predictors is ex- actly recovered, i.c., SIV (E) SVV (<) (exact recovery is typically not required in existing constrained analysis such as Woodworth et al. (017))Agawwa et al. (2018)). We show that if the number of sam- ples grow as - even the closest points on either side of the boundary of the set S (e) are correctly discriminated, which guarantees exact recovery of 8 (c). Once the exact set is recovered, beyond this we use standard learning theory tools to ensure R( (((*)== < R( S R(Φ*) + v.
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 standard result on ERM's sample complexity. Define a ++ such that o E arg mingeHo R(Φ)
& Ben-Davie 2014) For every > and 8 E (0,1). if He is a finite hypothesis class, Assumption 3 holds, and if the number of samples [D] is greater than 812 o log (2) then with probability at least 1 - 8, every solution Φi to ERM is an approxi- mation of expected risk minimization, i.e., R((++) < R(bt) R($+) + V.
Proposition 3. (Shalev-Shward
5

---

Published as a conference paper at ICLR 2021
Proposition 2 vs. vs. Since e, the sample complexity of EIRM grows least (maxx 22, bn }). Let us look at the two terms inside max- i) , growth term is similar to ERM, it ensures approx- imate optimality in the overall risk R. ii) / growth ensures the IRM penalty R' is less than A direct comparison of sample complexities in Propositions and suggests that the sample complex- ity 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.
3.3 OOD PERFORMANCE: ERM VS. IRM
We divide the comparisons based on distributional shift assumptions that decide whether ERM and IRM arrive at the same asymptotic solutions or not.
3.3.1 COVARIATE SHIFT
Assumption 4. Invariance w.r.t. all the features. Ve, E Eall and Ve E X, E|Y = x] X = z). Ve E Eail. X Pxx- and the support of Pe is equal to X.
As stated earlier, the first part of the above assumption follows from standard covariate shift assump- tions (Gretton et al. 2009) and is special case of the first part of the Assumption 1 with gh set to I. If कh = 1, then m (equation 2). which simplifies to m(x) = | X = xr). solves the OOD is is given as
problem equation [ A generative model that satisfies the above Assumption
y + 9(X)) + =", B[E"] = E" 1 xx, -[{-))]
(7)
In the above model X is the cause, Y is the effect, and g a general non-linear function (it satisfies Assumption 4 with = g). Next, we compare ERM and IRM's ability to learn under covariate shifts. In Figure a), we show Define = mine, A /0, |R($1) - R((2)), which measures
the minimum separation between the risks of any two distinct hypothesis in H.g
Proposition 4. Let be the square loss. For every v > 0, > 0 and (0, 1). if Hy is a finite hypothesis class, m E Ho, Assumptions 3 hold, and
if the number of samples D) is greater than max { 8L2 log( (|Hl ). 161 log(3)) then with probability at least 1 8, every solution & to EIRM (equation 6 satisfies R(m) S R(Φ)= s R(m) + v. If also v < R, then &- = m.
2 2|Hool 6 then with probability at least 1 1-6, 6, every solution o to ERM satisfies R(m)= R R( R(m) + If also < R, then Φi
if the number of samples |D] is greater than of
SL
log(
Implications of Proposition 4 ERM and EIRM both asymptotically achieve the ideal OOD solu- tion; the above proposition helps compare them in finite sample regime. The second term inside the max for EIRM, 101" ()). does not depend on the size of the hypothesis class. Hence, for large hypothesis classes, the sample complexity of EIRM equals
81,2 D² log( 4 Consequently, the
sample complexity of EIRM and ERM differs by a constant $1.² 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 tional covariate shift assumption in Proposition In Proposition we assumed square loss, but a similar result extends to cross-entropy loss as well.
Ø ), to Proposition [ Omax 27 A H); the addi-
4 helps get to a lower sample complexity of O( (,))).
3.3.2 DISTRIBUTIONAL SHIFT WITH CONFOUNDERS AND (OR) ANTI-CAUSAL VARIABLES
In this section, consider more general models than equation 7 which only contained cause X and effect Y 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.
Categorical (x (o0E),).),) Vo E Eir, T >0
y 7 (2i) + + 80 1 2i, Bloa] = 0, B[( = 0², 1&91 up
(8)
X + 5(2i, kg)

---

Published as a conference paper at ICLR 2021
We assume that Z1 component of S is invertible, i.e. 3 S such that S(S (z)) = Zi, and n # O. x nin
de E Etr, T5 2 Eerl 0. Define " = =|X* Ve E Err. 2 is positive definite, The support of distribution of 7e (7 }5), Pzz,, is bounded and the norm of S, |||||| a max (S) (maximum singular value of S), is also bounded.
In the above model, 7i is the cause of X and Ye but may not be directly observed. 75 may be arbitrarily correlated with Zi and E.. We observe a scrambled transformation of (Zi; 75) in X If Z is effect of Y (Z5 t Y + Nº), then Z5 is anti-causal variable. If He causes both 5ahh (E" t H + No) and Z (Z + H + Nº), then H is confounder. In both these cases, Zj is spuriously correlated with the label Y Consequently, the standard ERM based models estimate Z5 from X and end up being biased w.r.t the desired OOD model, which does not use
755- If Assumptions 2 hold, then linear model ST (Xe predict, 7 TT = =TZZi) ZT) solves the OOD problem in equation ( S, m 7 in Proposition [ and it relies only on Z In the supplement (Proposition 17), we prove that ERM based models do not recover STyy
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 E N if E > r r+ n/rr and for all non-zero I E R
dim
dim (sman (xxaa<xxx-------
BEE.. >
where span is the linear span, dim is the dimension, and recall n is dimension of X This assump- tion checks for diversity in the environments and holds almost everywhere (Arjovsky et al. 2019)
Assumption 7. Inductive bias. Hz is a finite set of linear models parametrized by + E RR" (output X*). STy Ha. 3 > 0, 0, s.t. V& E Ha. w S ||1b112 S 22 & 2w S ||S 7112 S 342/2 Ω
Informally stated, the above assumption requires the OOD optimal predictor ST to lie in the interior of the search space and not on the boundary. If Assumptions [ 7 hold, then Assump- tion holds. Hence, we use the bounds I. and T' on and (F(w (T.)
Ow 1==10 respectively in our next result. Define the minimum eigenvalue across all Σe as Amin = min Amin (Ee),
4th (21-16v2) -iin [E,,1 (wAmin)² and r
1
V //err Next, we analyze how EIRM learns STy.
3
Proposition 5. Let E be the square loss. For every E E (0, eth) and E (0, 1). if Asssupptions
5
(with 12.
7 shold and if the number of data points D) is greater than 161."
=² log (2 then with
a probability at least 8, every solution & to EIRM (equation satisfies & = (S 7)a, where -
a (1+1ve1--vaa
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 [ (Š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 S 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 's incquality.
Implications of Proposition 5 1. Convergence rate of ERM vs. EIRM: Recall that is the slack on IRM penalty R'. If e is sufficiently small and the data grows as ( every solution & to EIRM (equation 6 is in ve radius of the OOD solution, i.e., | ST
7|| = OVVE). We contrast these rates 2] v x. If the data grows as to ones in The covariate shift setting (Section 3.3.1 Let y
O (+zz ), then both ERM and EIRM solution converge to V as | T/| = (VV)) (from Proposition 4.. This shows that EIRM works in more settings (Proposition convergence rate of ERM.
than ERM while matching the
2. Sample complexity grows polynomially in data dimension to ensure OOD generalization:
Next, we set e = j&tt with μ E [0, 1). and |Eer = 2n (satisfies Assumption 6 f r 1). A simple manipulation of terms in Proposition 55) shows that sample complexity with quadratic growth in data
dimension n, log ( (2[Hel) )). ensures & = (ŠTz)a with a E [[+v/(VZ-1) 1-VR(v2-1)+
7

---

Published as a conference paper at ICLR 2021
3. Comparison with Proposition 2: Lastly, we contrast sample complexity of EIRM in Propo-
sition 5 O ), to Proposition 2
(max{ y2 }); the additional distributional assumptions in
Proposition
5 help arrive at a lower sample complexity of ( The bound in Proposition 2
(max + pz }), is larger than the one in Proposition O ((tz)) and Proposition
O +z) but is
more general as agnostic to the distributional assumptions.
A simple illustration summarizing Propositions 4
5
Set S to identity in Assumption 5 Recall
Zi and Z5 from Assumption 5 Since S is identity X can be written as Xf. X5), where X Zi
and X Z5. If X5 1 E.. then [X is invariant and Assumption 4 holds. This corresponds
to the setup in Figure Πa). 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 X If + E" Nº. then this corresponds to the setup Figure(Tb), X is the cause and X; is spuriously correlated with label ye through the confounder ex If X; Y*+NN, then this corresponds the setup in Figure [ Xf is the cause and X3 is anti-causally related to the label yE. In both these cases, the ideal OOD solution that solves equation Twill only exploit Xf to make predictions. From Propositions 5 it follows that IRM when fed with ( samples, it is in VE radius of the target OOD solution, while ERM is asymptotically biased and exploits X5 (Proposition 17).
We define polynomial version of the model in Assumption 5 We only need to change y <- y (Z1) + E" to Y + 72f) + Eo is a polynomial feature map of degree p defined as Sp : Re - Tee where e is the dimension of the input Z1, (,WW) (W, W Ø W. W...p times @ W)) ((W and @ is the Kronecker product. Also, c = Ee c². Can we directly use the analysis from the linear case by transforming X 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 X and Z We present the main result informally below (details are in the supplement).
Proposition 6. (Informal statement) For sufficiently small e and ā E (0, 1), if ifssumptionss similar to Propositionn holld and DI 2 164" log( 2(Hlll, ), then with a probability at least 1 1- 6, every solution t to EIRM (equation 6 satisfies :- = ST y(a), where STay is the OOD optimal model (defined in the supplement), a E I+TVE --vll
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 [ are now required to grow as Onz)). As a result, in the sample complexity analysis we discussed, we replace n with re" to obtain that a sample complexity of O log (2))te)) ensures & = 3m x(a) with 1
a E 1+vp(v2- 1)' 1-))(v2-1)) In the infinite hypothesis case, the main change in the results is that we replace |Ho] with an appropriate model complexity metric (Shalev-Shwartz & Ben-David 2014). Consider Proposition 5 sample complexity of ( log (a) ensures 1 (ST>)) with a E showed the benefits of IRM for polynomial models and other extensions (non-linear S) are future work. In the supplement, we provide dialogue explaining how our work fits in the big picture.
l (I+V(((2-1) ( -VVVI-1) in contrast to O nn log (2)%)l)) in the finite hypothesis case. We
4 EXPERIMENTS
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. (019). We divide the training data MNIST digits into two environments (e = 2, equally and the testing data in MNIST digits is assigned to another environment (e 3). Xe: gray scale image of the digit, Ye: label of the gray scale digit (digits 2 have g = and digits < 5 have Y = 0). X final colored image and Y": final label are generated as follows. Define Bernoulli variables G. N. No that take a value 1 with probability 8, B and 3 and 0 otherwise. Define color variable C", where C = 0 red and C 1 is green. Let @ denotes xor operation.
Y + L(Xg), L: Human labeling,
y + L(Xg) @ N, Corrupt the original labels with noise
C + G(Y @ No) + (1 - GN @ N Use G to select b/w anti-causal or confounded (10)
X + T(X,, C transformation to color the image
8

---

Published as a conference paper at ICLR 2021
If the probability 8 = 1, then G 1 and that gives back the original CMNIST in (Arjovsky et al. 2019), which we call anti-causal CMNIST (AC-CMNIST). If f = 0. then we get confounded If 0 < 8 < 1, we get hybrid dataset (HB- (HB-CMMIIST). The above model in equation 10 has features of the model in Assumption 5 where X C L, T take the role of 7 75. 7, S. We set the noise N parameter B = 0.25, and the parameter for Ne in = [0.1, 0.2, 0.9). Color is spuriously correlated with the label; P(C = 1|Y 1) varies drastically in the three environments (00.9,.85 0.1] for AC-CMNIST), In the variants of CMNIST we discussed. P(Y varies across the environments. We now define a covariate shift based CMNIST (CS- (CS-CMNIST) (P(Y"*IX") induce spurious correlations. Generate a color C uniformly at random. Select the pair (X go C*) with probability If the pair is selected, then color the image X < (X
colored MNIST
the three environments 13" _8]
invariant). We use selection bias to
v/ if the label 9 and C' are the same, else select them with probability vlf. g. C*) and Y < Y Selection probability g for the three environments are [&b v² e = 0.2, 0.9]. Due to the selection bias, color is spuriously correlated, P(C 1 1) varies drastically [0.9, 0.8, 0.1]. We provide the graphical models for the CMNIST variants and computations of (P
P(Y*(X)) in the supplement.
FRM
RM
Tas
Number sampes
N.Mmmrr samples
onor
Number of samplies
0) Number samples
Figure 2: Comparisons: a) AS-CCNIST, b) CF-CMNIST, c) AC-CMNIST and d) HB-CMNIST.
4.1 RESULTS
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 (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 2h (as predicted from Proposition 4 IRM outperforms ERM in the remaining three datasets (as predicted from Proposition 5) as seen in Figure 2p.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 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).
CONCLUSION
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 (ooossyyysky 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 model complexity, number of environments on the sample complexity and its distance from the OOD solution in distribution shifts that go beyond covariate shifts.

---

Published as a conference paper at ICLR 2021
6 ACKNOWLEDGEMENTS
This work was supported part by the Rensselaer- AI Research Collaboration (part of the IBM AI Horizons Network).
REFERENCES
Alekh Agarwal, Alina Beygelzimer, Miroslav Dudik, John Langford, and Hanna Wallach. A reduc- tions approach to fair classification. arXiv preprint
2018.
Kartik Ahuja, Karthikeyan Shanmugam, Kush Varshney, and Amit Dhurandhar. Invariant risk min- imization game. In International Conference on Machine Learning, 2020.
Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, and Mario Marchand. Domain-adversarial neural networks. arXiv preprint arXiv:1412.4446: 2014.
Isabela Albuquerque. João Monteiro, Tiago Falk, and Ioannis Mitliagkas. Adversarial target- invariant representation learning for domain generalization. arXiv preprint 10:iv:/911.00000 2019.
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint wrXiv:/907.029500 2019.
Robert Ash, Robert, Catherine A Doleans-Dade, and Catherine. Probability and measure theory. Academic Press, 2000.
Sara Beery, Grant Van Horn, and Pietro Perona. Recognition in terra incognita. In Proceedings of the European Conference on Computer Vision, pp. 456-473, 2018.
Shai Ben-David and Ruth Urner. On the hardness of domain adaptation and the utility of unla- beled target samples. In International Conference on Algorithmic Learning Theory, pp. 139-153. Springer, 2012.
Shai Ben- David, John Blitzer, Koby Crammer, and Fernando Pereira. Analysis of representations for domain adaptation. In Advances in neural information processing systems, pp. 137-144, 2007.
Shai Ben-David, John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jennifer Wort- man Vaughan. A theory of learning from different domains. Machine learning, 1511-2): 2010.
Shiyu Chang, Yang Zhang, Mo Yu, and Tommi S Jaakkola. Invariant rationalization. In International Conference on Machine Learning, 2020.
Shai Ben David, Tyler Lu, Teresa Luu, and Dávid Pal Impossibility theorems for domain adaptation. In Proceedings of the Thirteenth International Conference Artificial Intelligence and Statisties, pp. 136, 2010.
Alex J DeGrave, Joseph D Janizek, and Su-In Lee. Ai for radiographic covid-19 detection selects shortcuts over signal. medRxiv, 2020.
Zhun Deng, Frances Ding, Cynthia Dwork, Rachel Hong, Giovanni Parmigiani, Prasad Patil, and Pragya Sur. Representation via representations: Domain generalization via adversarially learned invariant representations. arXiv preprint arXiv:2006. 14788, 2020.
Arthur Gretton, Alex Smola, Jiayuan Huang, Marcel Schmitifull, Karsten Borgwardt, and Bernhard Schölkopf. Covariate shift by kernel mean matching. Dataset shift in machine learning, 3(4):5, 2009.
Ishaan Guirajani and David Lopez-Paz. In search of lost domain generalization. arXiv preprint arX 2020.
Masanori Koyama and Shoichiro Yamaguchi. Out-of distri-iistiitttion generalization with maximal in- variant predictor, arXiv preprint arXiv:2008.0/88 2020.
10

---

Published as a conference paper at ICLR 2021
David Krueger, Ethan Caballero, Priol, and Aaron Courville. Out-of-distee distriiittiittionn generalization via risk extrapolation (rex). arXiv preprint XXi1:20033006688 2020.
Jacobsen, Amy Zhang, Jonathan Binas, Remi Le
Alan J Laub. Matrix analysis for scientists and engineers, volume 91. Siam, 2005.
Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M Hospedales, Deeper, broader and artier domain generalization. Proceedings of the IEEE international conference on computer vision, pp- 5542-55500 2017.
Divyat Mahajan, Shruti Tople, and Amit Sharma. Domain generalization using causal matching. arXiv preprint arXiv:2006.07555 07500, 2020.
Toshihiko Matsuura and Tatsuya Harada. Domain generalization using mixture of multiple latent domains. In AAAI, pp. 11749-11766 2020.
Krikamol Muandet, David Balduzzi, and Bernhard Schölkopf. Domain generalization via invariant feature representation. In International Conference on Machine Learning, pp. 10-18, 2013.
Artidoro Pagnoni, Stefan Gramatovici, and Samuel Liu. Pac learning guarantees under covariate shift. arXiv preprint arXiv:1812.06394 2018.
Judea Pearl. Causality. Cambridge university press, 2009.
Vihari Piratla, Praneeth Netrapalli, and Sunita Sarawagi. Efficient domain generalization via common-specific low-rank decomposition. In International Conference on Machine Learning. 2020.
levgen Redko, Emilie Morvant, Amaury Habrard, Mare Sebban, and Younès Bennani. Advances in Domain Adaptation Theory. Elsevier, 2019.
Shai Shalev-Shwartz and Shai Ben-David. Understanding machine learning: From theory algo- rithms. Cambridge university press, 2014.
Damien Teney, Ehsan Abbasnejad, and Anton den Hengel. Unshuffting data for improved generalization. arXiv preprint arXiv:2002.11894, 2020.
Blake Woodworth, Suriya Gunasekar, Mesrob Ohannessian, and Nathan Srebro. Learning non- discriminatory predictors. urXiv preprint arXiv: 2017.
Keyulu Xu, Mozhi Zhang, Jingling Li, Simon Shaolei Du, Ken-Ichi Kawarabayashi, and Stefanie Jegelka. How neural networks extrapolate: From feedforward to graph neural networks. In International Conference on Learning Representations, 2021. URL htthttts////ppeeeeeees
net//orrum?ic-Ol-cnne
Han Zhao, Remi Tachet des Combes, Kun Zhang, and Geoffrey J Gordon. On learning invariant representation for domain adaptation. arXiv preprint arXiv: 1901: 09453,,993,,... 2019.
SUPPLEMENT
7.1 DIALOGUE ON IRM CONTINUED
[In the original IRM paper by Arjovsky et al. (2019), we left two graduate students ERIC and IRMA strolling along the Parisian streets on warm summer evening. lot has transpired since then. ERIC is now holed up inside his room at Cité Universitaire and IRMA has returned to her parents' home Provence. The two are talking through Zoom call.)
After reading Arjovsky et al. (2019) and Guirajani & Lopez-Paz 2020), was not sure of how to understand the settings where IRM is beneficial over ERM and vice-versa? From Gulrajani & Lopez-Pac (2020). I understood that ERM continues to be the state-of- the-art if sufficient care is taken in performing model selection.
ERIC:

---

Published as a conference paper at ICLR 2021
IRMA: But after reading this manuscript understand that there 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 's audio and video freeze for minute, leaving ERIC to ponder this conundrum for little bit on his own. IRMA stops her video and continues only with audio and the conversation able to resume.]
IRMA: Sorry, how much did you hear?
ERIC: You were just starting to explain why ERM can be at disadvantage, but didn't hear anything after that.
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-Pazz (2020) and will come to it in a bit.
IRMA:
Yes, you are right! The datasets are human labeled images if recall correctly. In these datasets it's safe to assume P(Y+|X*)) is close to invariant across domains (as long as the cohort of humans the different domains are not very different). What happens when P) varies lot?
ERIC:
Yes, when PP({))*)) varies a lot, IRM can be at an advantage over ERM. In this manuscript, I learned that when data generation involves confounders or anti-causa vari- ables, it is possible to show that IRM methods converge to desirable OOD solutions with good convergence rates.
IRMA:
ERIC: Interesting! should read that. The experiment on colored MNIST by Arjovsky et al. (019)) 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 confounder and anti-causal). IRM retains advantage on all but the first dataset. 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 better than ERM and vice-versa?
IRMA: Yes, actually have been thinking about this problem and would be happy share my initial thoughts. Identifying representation d that leads to invariant conditional dis- tributions 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 through x to get the 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 come to, why IRM did outperform ERM Guirajaní & Lopez-Par 2020
these representations
ERIC: How about when there is complete overlap?
Yes, if there is strong complete overlap in the representations from the two domains, then is possible that IRM can help. In fact I believe in such settings, the ERM models trained on the two domains disagree lot, then that can be strong indication towards the models heavily exploiting spurious correlations. In such cases, believe IRM again offer some advantage.
IRMA:
12

---

Published as a conference paper at ICLR 2021
ERIC: I will try to put these ideas down paper and meet you again for a discussion.
IRMA: Great! truly hope that it can be in person at the café in Palais-Royal where we first started this conversation. So long for now!
(The students end their call.)
7.2
SUPPLEMENTARY MATERIALS FOR EXPERIMENTS
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 httph:t////// /ttt:s://gthus/ wwt. com/IBM/OoD
7.2.1 1
CLASSIFICATION
We first describe the model and other training details.
Choice of His 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.001 1 is used for each layer. We use a learning rate of 4.9e-4, hatch 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 large penalty 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).
A. Covariate shift based for CMNIST
We provide the generative model for CS-CMNIST below.
Y L(Xg), O Unifurm((O0)))
U + Bernoulli (C @ y9)00 + +(1 - (C" + +Y,))
(II)
X + T(X*, 1= + U-
A.1. Compute (CC and PYY
We compute P(C and P(Y for the covariate shift based CMNIST described by equa-
tion 11 ((C))))) helps understand how the spurious correlations vary across the environments. P(Y Xxx) helps us understand if the covariate shift condition satisfied or not. Compute the prob-
ability PCC Y 1) P 1, U 1) follows.
P(C = 1, = 11Ue =1)
PPCC
i 1,Ue =1)= P(C = 1, 1 = 1) PC 0, Ye 1 = 1)
P(C = 1, Ye = =110= = 1) = P(C =1, = -1,0=
(1 - (1-
Σa, P(C a, Yú = b,00 1)
2
(12)
PPC = 0, ye 1 = 1) = P(Ce = = =
2 e
Ea,b P(C a, Yg = -
PCC 1|Y = 1, U = 1) (1 (*)
From the above simplification we gather that PC = 1 = 1) is 0.9, 0.8 and 0.1 in the three environments. Define Pt P(Ye 1 In MNIST data it is reasonable assume deterministic labeling, i.e. Pr = or = 6. From the way the data is constructed we can assume that T is invertible, i.e. from colored image X can get back the grayscale image and the color Xe C g>
We now move to computing P(Y=)X))). We assume ye in the simplifcation below.
13

---

Published as a conference paper at ICLR 2021
X
Y
Y
U
X
C
X
grayscale image
U selection variable
y digit label
environment index
C
color index
X colored image
y final label
Figure 3: Graphical model for CS-CMNIST
-0 0)P(U* P(Y = == = 0 = 1|X,) = P(Y = = -)PU0 = 11 = 1, = 0)
= -0.5pree
P(Y - = 0, U 1|X;) P(Y P|x3) = 0
1 -, 0)
-0.((1 (1 1h))
P-Y = 11x = P.Y g = 1|,;
9: C = 0 =)
(13)
P(Y = 1, C -0,U- = =
F(Yg 10 0, U 1(x)) + +(Y O 0,0' = 1[x;)
pit'c
IPY = 11x*) 2ptibe
pt
Ve
Observe that under standard assumption of deterministic labeling Pl = 1, P(Y = 1(x") remains invariant and for high values of pt it is relatively stable across the environments.
A.2 Graphical model for covariate shift based CMNIST. In Figure model for covariate shift based CMNIST described in equation II
we provide the graphical
A.3. Results with numerical values and standard errors. Table 2 we provide the numerical values for the results showed in the Figure a) along with the standard errors.
B. Colored MNIST with anti-causal variables (AC-CMNIST)
B.I. Compute P(C)) and P(** We first compute P P(C 1 = P(Y @ No 1 = 1) = P(N" = ojУ = 1) = 8 Therefore, P and 0.1 in environments 2. and 3 respectively. Next, compute PYY We assume Y in the simplfication below. We also assume deteministic labeling. In the simplfication that follows
= 1y" = 0.9. 0.8
use 3 = 0.25.
14

---

Published as a conference paper at ICLR 2021
Method
Number of samples
Test error
ERM
1000
47.27 + 0.72
IRM
1000
42.84 + 0.94
ERM
5000
23.91 1.23
IRM
5000
28.85 + 5.21
ERM
10000
15.80 + 0.63
IRM
10000
14.80 + 0.30
ERM
30000
9.86 + 0.32
IRM
30000
11.20 = 0.64
ERM
60000
8.15 + 0.37
IRM
60000
9.62 + 0.67
Table 2: Comparison of ERM vs IRM: CS-CMNIST
P(Y 1, C 0/Xg)
P(Y 1, C 0|
PI( 111 1)P(C 0|r 1) 0.75,
P(Y 0, C 0|x3g)
P(Y 1, C 0|x35) IP( 011 1)P((" = 0|r 0) 0.25 -8)
3
(14)
P(Y 1[X*) = P(Y 1 1x C = 0) 2,3 +1
P(Y 1(X*) = 0.25 (For environment 1)
P(Y 1[X*) 0.428 (For environment 2)
P(Y = 1(X") = 0.964 (For environment 2)
B.2 Graphical model for anti-causal CMNIST In Figure 4 we provide the graphical model for AC-CMNIST described in equation TO (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 along with the standard errors.
C. Colored MNIST with confounded variables (CF-CMNIST)
C.1. Compute PCC and P(Y.)."). We start by computing P(C Recall that C
(N @ No). In the simplfication that follows we use 3 0.25.
P(N P(N 0 1)
0.75 0.5
P(N= = 0|8* = 1|
0.75
0.75 * 0.5 + 0.25 * 0.5
(15)
We use the above to compute 0 -N -
Y 1) (1 8 0.255 = 0.75 - 0.5. For environment 1, the above probability P(C = 0|Y 1) is 0.7, 0.65 and 0.30 respectively. Next, we compute the probability P(Y*)x*).
Suppose Yı 1 for the calculation below. We also assume deterministic labeling.
15

---

Published as a conference paper at ICLR 2021
Method
Number of samples
Test error
ERM
1000
78.04 0.70
IRM
1000
77.12 + 1.00
ERM
5000
75.23 1.04
IRM
5000
74.89 = 1.08
ERM
10000
74.68 + 1.23
IRM
10000
60.42 + 1.72
ERM
30000
69.87 = 0.39
IRM
30000
49.61 + 2.57
ERM
60000
71.96 = 0.55
IRM
60000
32.78 2.70
Table 3: Comparison of ERM vs IRM: AC-CMNIST
P(Y 1, C 1|x9)
P(Y = 1, C, 1|x,)
P(y== = 1)P((C- = 1| 0) = 0.75g
P(Y 0, C =1
P(Y" = 1,C" 1
P(Y = 1)TP(C = 1/N 1) = 0.25 - B²)
333e
(16)
P(Y = 1|X*) P(Y 1 1 C =1= 23 +1
P(Y 1(X*) 0.25 (For environment 1)
P(Y = ) = 0.428 (For environment 2)
P(Y = 1 0.964 (For environment 2)
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 we provide the numerical values for the results showed in the Figure bb) along with the standard errors.
D. Colored MNIST with anti-causal variables and confounded variables (HB-CMNIST)
D.I. Compute P(C and P(Y
We start by computing P(C" Recall that = = -C G= with probability 0 and 0 otherwise.
P(C = 1,y
PPC =
(17)
PCC =) IP 0, 0,
P(C = 1, - IP(C" = 0, = - = (0.75 We used 0 0.8 in the experiments. P(Y*[X*) can be computed on the same lines shown for anti-causal and confounded model and varies significantly across the environments.
(0.25 +
(18)
(19)
D.2 Graphical model for confounded CMNIST.
In Figure 4 we provide the graphical model for confounded CMNIST described in equation 10 (for
0 < 0
<
D.3. Results with numerical values and standard errors. Table 5 we provide the numerical values for the results showed in the Figure d) along with the standard errors.
16

---

Published as a conference paper at ICLR 2021
Method
Number of samples
Test error
ERM
1000
65.21 0.64
IRM
1000
65.60 + 0.53
ERM
5000
67.91 0.40
IRM
5000
67.59 + 1.14
ERM
10000
66.92 = 0.30
IRM
10000
59.26 + 1.62
ERM
30000
67.32 = 0.28
IRM
30000
47.01 + 1.37
ERM
60000
67.62 = 0.31
IRM
60000
44.92 0.84
Table 4: Comparison of ERM vs IRM: CF-CMNIST
Method
Number of samples
Test error
ERM
1000
65.38 + 0.56
IRM
1000
66.98 + 0.61
ERM
5000
62.92 + 0.99
IRM
5000
61.09 + 1.74
ERM
10000
62.08 + 1.19
IRM
10000
51.46 1.22
ERM
30000
60.48 + 0.42
IRM
30000
34.96 1.47
ERM
60000
59.70 + 0.72
IRM
60000
32.80 + 0.55
Table 5: Comparison ERM vs IRM: HB-CMNIST
7.2.2 REGRESSION
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 written as
H + N(0,
Xi + o2(0) W
(20)
H ye Wi-g Xi -N Ho
X w,-2Y y 1,), W
H is the hidden confounder, X [Xf X the observed covariate vector, Y is the label. Dif- ferent 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 describe them below. de is environment
17

---

Published as a conference paper at ICLR 2021
CF-CMNIST
AC-CMNIST
HB-CMNIST
X
Y
X
Y
Y
N
N
Y
N
Y
X
X
C
X grayscale image
No environment dependent moise
6
noise
digit label
environment index
Y final label
C
color index
X
colored image
G select b/w anti-cousal
confonnded
Figure 4: Graphical model for CF-CMNIST, AC-CMNIST, HB-CMNIST
H
H
Ife
17"
Xi
X
Xf
Y
XG
X
ye
X
Xf
Y
X5
CS Rogression
Re-Reerrssoon
AC-Regresion
HB- Regression
Figure 5: Graphical models for the different regression datasets
dependent standard deviation, we use two training environments with o = 0.2 and respectively and both environments.
2.0
Covariate shift case (CS-regresion): In this case, fix Wh-22, Wh-+y- and W Wy-2 to and draw each entry of W1-+ from (0, 1) and set Wh-1 to identity.
Confounded variable case (CF-regression): Set Wy-12 to zero and we draw each entry in W1-ty, Wh++1, WR-+2. Wi-+y from (0, 1).
Anti-causal variable (AC-regressioon): Set Wh-y. W and WR-22 to and draw each entry of W1-y and W y2 from (W( 1).
Hybrid confounded and anti-causal variable case (HB-regressioon): Draw each Wh-y.Wh--1 Wh++2. W1-vy and W Wy-+ from (0,
We also present the graphical models for the four types of models used in Figure
7.2.3 CHOICE OF Hi AND OTHER TRAINING DETAILS
We use linear model that takes input X 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, le 5, le - 4, le 3, 1e le (with 4:1 nainvvaliaation split). We average the results over 25 trials.
18

---

Published as a conference paper at ICLR 2021
7.2.4 RESULTS
We discuss results for the case when the length of the covariate vector X is 10. The desired optimal invariant predictor is W = [Wi-ty:0 0j. 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*ll2. In Figures 6 & 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 in the covariate shift case (See Figure 6) there is clear winner between the two approaches. There are gains from using IRM in the other cases (Figures 8 99. 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 ERM is also very small. In addition to the Figures
9 we provide
the tables (see Table 6
with the numerical values for the mean model estimation (and the standard error) error shown in the figures.
8 9
Model estimation error
1.75
ERM
Model estimation error
1.75
ERM
1.50
IRM
1.50
IRM
1.25
1.25
0.75
0.75
0.50
0.50
0.00
0.00
500
1000
1500
2000
2500
3000
3500
4000
500
1000
1500
2000
2500
3000
3500
4000
Number of samples
Number of samples
Figure 6: Comparisons: n = 10 CS-regression
Figure 7: Comparisons: n = 10 CF-regression
Model estimation error
1.75
ERM
Model estimation error
1.75
ERM
1.50
IRM
1.50
IRM
1.25
1.25
1.00
1.00
0.50
0.50
0.25
0.25
1000
2000
2500
3000
4000
1000
1500
2000
2500
3000
3500
4000
Number of samples
Number of samples
Figure 8: Comparisons: 10 AC-regression
Figure 9: Comparisons: 10 HB-regression
19

---

Published as a conference paper at ICLR 2021
Method
Number of samples
Model estimation error
FRM
0.65 + 0.06
IRM
50
0.57 + 0.10
ERM
200
0.18 + 0.02
IRM
200
0.13 + 0.02
ERM
500
0.08 + 0.005
IRM
500
0.08 + 0.014
ERM
1000
0.037 + 0.004
IRM
1000
0.051 + 0.007
ERM
1500
0.021 + 0.002
IRM
1500
0.045 + 0.007
ERM
2000
0.018 + 0.002
IRM
2000
0.035 + 0.006
Table 6: Comparison of ERM vs IRM: 10 CS-regression
Method Number of samples
Model estimation error
ERM
50
1.48 + 0.15
IRM
50
0.59 + 0.10
ERM
200
0.30 + 0.03
IRM
200
0.20 + 0.03
ERM
500
0.14 + 0.01
IRM
500
0.15 + 0.02
ERM
1000
0.10 + 0.02
IRM
1000
0.09 + 0.01
ERM
1500
0.07 + 0.01
IRM
1500
0.10 + 0.01
ERM
2000
0.07 0.01
IRM
2000
0.07 0.01
Table 7: Comparison of ERM IRM: 10 CF-regression
7.3 PROOFS FOR THE PROPOSITIONS
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 fg fI,... ..... be a sequence of i.idd random variables and assume that for all i, E[0;] p and Pla < 0 < bj
Then, for any 0
< 2 exxp -2 me?
8
11 c
(b a)²
We restate the propositions from the main body. Next, we prove Proposition T from the main body of the manuscript.
20

---

Published as a conference paper at ICLR 2021
Method
Number of samples
Model estimation error
FRM
1.57 + 0.18
IRM
50
0.82 + 0.20
ERM
200
0.90 + 0.08
IRM
200
0.63 + 0.08
ERM
500
0.81 + 0.08
IRM
500
0.61 + 0.05
ERM
1000
0.75 + 0.07
IRM
1000
0.53 + 0.04
ERM
1500
0.79 + 0.07
IRM
1500
0.53 + 0.05
ERM
2000
0.77 + 0.08
IRM
2000
0.54 + 0.04
Table 8: Comparison of ERM vs IRM: n 10 AC-regression
Method Number of samples
Model estimation error
ERM
50
1.81 + 0.19
IRM
50
1.13 + 0.18
ERM
200
0.87 + 0.09
IRM
200
0.68 + 0.10
ERM
500
0.83 + 0.09
IRM
500
0.67 + 0.07
ERM
1000
0.76 + 0.07
IRM
1000
0.61 + 0.05
ERM
1500
0.77 + 0.07
IRM
1500
0.63 + 0.05
ERM
2000
0.74 + 0.07
IRM
2000
0.58 +:: 0.05
Table 9: Comparison of ERM vs IRM: 10 HB HB-regression
Proposition 7. If f is square loss, and Assumptions 72 hold, then t. solves the 0OD problem (equation 7
Proof. Define a predictor w o Φ", where Ø" is defined in Assumption T Let us simplify the expres- sion for the risk for this predictor R. (w (²)) using square loss for f. Recall Ze +x(X))
21

---

Published as a conference paper at ICLR 2021
( (w o Φ6) Eº - E |-] + E [ z-] - (w o x (xx) 21
|
R
(Ye E [ye Ex [x<] 13c (B [r<] | (w x (X +
2
(ve - B (va zz0 ( 1A | (m) o (x(()
E
(Y m m(29))²
- (m(z*) w(29))² 2
2E (Y -)(=m)) (m(2))(m((2) - w (29))
R (ve-m ml (29)) )
(m(29) w/ 2z)))
= E (m((z) - 2(9))²²
(21)
In the above simplification in equation 21 rely on the law of total expectation.
we use the following equation 22 and equation 23 which
E.
( (r ml m(Z²))
x(()) -m())))))()>))
(m(Z")
EE"
E [(r- ml ( (m(z") - w(z))
LE E.
(m(z²) w0 (20)
=0
(22)
E
(yr- m m(z))
27
E
Eey (v= m((²))
z3|
E Var [y<]za]
20]
= g²
(23)
In the last equality in equation 23
we use the Assumption and obtain g?. Therefore, substituting in equation 21 achieves risk of E² for all the environments.
Ve E Eall, He [m
(24)
Let us simplify the expression for the risk achieved by every predictor f E F in the environment 9 following the steps similar to equation 21
Consider the environment g that satisfies Assumption 2
(f) = EE
(Y" 19
27
E7 (ETY"9229) - ((x)) 27
(Y" - =2(17)9))) (FTT)))))) -(x))]
E9
+ 2
E2
1 - m222))² 27
+ 1E" (m(z9) - =(X9)) )²1
+ 2 (x--
(x*=m(29)) m(29))
(m(22)-
(y" - m (20) "]
(m(29) s(x9)) "]
E2
+ EE9
= + E0/ (
1
- -((x)) 27
(25)
In the above simplification in equation 25 we use the following equation 26
Ea (Y" 1y--m(f (((=)) (m(Z²) -1(x)) Ed E [y- - m(Z4) (m22) f(1))
|
= [ - mm(2)) (m(29) (
(11129)-(((9))) f(x9)) (m(29)-f(x9))
= 0 (Last equality follows from the Assumption 2
(26)
22

---

Published as a conference paper at ICLR 2021
Therefore, for environment satisfying aasumption2ffrom equation 25 R (f) 2 5². Therefore, we can write that
it follows that for all f E F,
of
max (f) (f 2 eEEnt
(27)
Therefore, from equation 27
directly follows that
min max ()) FEF eaEal
(28)
We showed in equation 24 Rd (m 2-) c² for all the environments.
Hence, f = m o 4* 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 build these propositions.
In Proposition 2 our goal is carry out 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) ii) unlike ERM, IRM 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 separable fashion. To deal with ii), we define a parameter i that measures the minimum separation between IRM penalty and c. We show that as long as this separation for all the predictors in the hypothesis class is positive, then we can rely on x-representative property from Shalev- Shwartz & Ben-D (0114) applied to the new estimator that we build to show that the set of empirical invariant predictors SIV (c) are the same exact invariant predictors SIV(e).
In Proposition 5 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|Arooskky 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 be in the ve 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 [S] 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 particu- lar, we focus on parametric model families that are Lipschitz continuous. The extension to infinite hypothesis classes based on carefully exploiting the covering number based techniques Shharrv-Swwatz & Ben-David (2014) for the IRM penalty estimator that we in- troduced. 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
V Rf (w- ..)) +) OE [[(w $(X'), Y
al(ww ((X*), Yc)
E
Ow
Ow
1=110
and
(JE [[(w ((X+), Yo)
2
de(w (w ((X)), Ye)
2
1.0 (u-
on im
Ow
w=1.0/
=
Ow
w=1.0
(29)
23

---

Published as a conference paper at ICLR 2021
In the above simplification, used Leibniz integral rule and take the derivative inside the expecta- tion.
Also can write [[X]² B[AB], where A and B are independent and identical random variables with same distribution as X. Therefore, we consider two independent data points (X", Y Ie and (Xe, ye)= ~ Pe
V w= (w- 2)1l= =E= [ ((.)) Y
(ollw +(X"), Y
(30)
Ow
1=110
Ow
In the above the expectation E taken over the joint distribution over pairs of distributions of pairs (Xe, y²), (Xe, yo) from the same environment e.
We write
R (Φ) = E v1
+)11
eEErr
(31)
(Ōelw x(X"), Y
de(w ((X"), 89
T-
Ow
w=10
Ow
In the above simplification, used equation
30
Define joint distribution P over the tuple
(e, (X", Y²), (X', )))), where
(X Yo) F and (X", Y*) ~ PP Also,
i ((ee, (Y), (X9, r (x" """,""""""" 8a)
(32)
We rewrite the above expression equation in terms of an expectation w.r.t P. which we represent as E follows
R ( E (DRlu: ((x)), ve)
af(w- $(
(33)
w=1.0/
Iw
Define
al(w o(X"), Y
all(w- +(X"), 870
(34)
Ow
1=110
Ow
1==10
Substitute equation
34 in equation 33 to obtain
R' ( 16 e' ("x
(91)]
(35)
We construct simple estimator A' (b) by pairing the data points in each environment. For simplic- ity assume that each environment has even number of points. In environment e, which has ne points construct 2 pairs. Define of such pairs
D = (108)-11 (1 vil Vi
(36)
2
R' ④
h, Y (X {2))
(37)
D
EEE
im1
There be other estimators of R() where separately estimate each term | 1.0 R (w. and 7" in the summation. We rely on the above estimator equation 37 as its separability allows standard concentration inequalities, e.g., Hoeffding's inequality.
24

---

Published as a conference paper at ICLR 2021
7.3.2 E-REPRESENTATIVE TRAINING SET FOR R AND R
We the definition of e-representative sample from Shalev-Shwartz & Ben-David (0144) and state it appropriately for both and R. Definition 1. A training set called e-representative (w.r.. domain 2, hypothesis H. loss E and distribution D) if
Wh H. [[(h))(hh))))))))))) ([()) S where R(h) = Ex(e(((X), Y and (X D.
(38)
Following the above definition, we apply the set of points δ defined above in equation 36
called e-representative w.r.t. domain X, hypothesis H4. loss e' (equation 33 and distribution (equation 32 if
v Ho, Ho,H,R (+) R' R())) Sc
(39)
11' (Φ) = B (1. (9, {-1))]
where
(from equation
34
and
(c, (X+, VC), (o)) ~ P.
Recall the definition of min |R' (Φ) e/.. Next, show that if Ò is 2 --rereesentative w.r.t X, Ha. loss é' distribution then the set of invariant predictors in equation (()V(()) (e) and the of invariant predictors in equation (5 (e)) are equal.
Lemma 2. If & > and H is & s-representative w.r. X. Hp. loss and distribution P. then (e)(e) (t)(
Proof. First we show (e)(e) S (c). From the definition of min |R (Φ) cl follows that VD Hx
R((4) 2 K
R ( (c) SE-N
(40)
Consider any in (c))
(
(41)
Given the definition of and equation 40 we obtain
R' RR(@) <c =
<c
(42)
Therefore, (E) C (*)(E-k). Also, it follows from the definition of the set (c) that S (e k) C (.V(())) Hence,
(eV(e)= ( (k)
(43)
Consider any SIV (e)
R' ( S e (From equation 43
A' (0) - R ( + R 5
(44)
R' ( S -
R'())]
From the definition of s-reppeeeentativeness it follows that R ($) R' (.) s and substituting this in equation we get
R' ( x/2 =( SNv(e) SVV(e)
(45)
Next we show (IV(e) _ (iV.)). Consider & E SIV
25

---

Published as a conference paper at ICLR 2021
( < A' (b) - R (0) ( Sc
(46)
S ( (((
From the definition of $ s-representattivensss it follows that R((D) R (0) s & and substituting this in equation 46 we get
R' S ct3 2
(47)
From equation 40 it follows that R (Φ) S c = ( S.. Therefore, E (c). This proves the second part (!V(e) C (e) and completes the proof.
Lemma 3. Ifk 0, represennaative w.r.t X, Ha, loss l and distribution PP (joint distribution
over (e, (e))) Y*) defined in Section 3.1 and and D is 5-representative w.r.t X, Ha, loss E and distribution P, then every solution to EIRM (equation 6 satisfies & is in (c) and R(b*) S
R(Φ)= < R(2²) + e, where 4 is the solution of IRM in equation 5
Proof. Given the condition in the above lemma, we able use the previous Lemma to deduce that 5tV (c) S (c). This makes the set of predictors satisfying the constraints in EIRM equation and IRM equation 5 the same.
T" solves equation 5 and क solves equation 6 From -rrepresentativeness we know that R( < R( (). From the optimality of & we know that R( ( S A((T)) (0* E 5ev StV (e)). Moreover, from v follows.
2 g-representativennss we know that R((*) I R(**)= v. 2-
We combine these conditions as
R( - 2 S ((()) < (x))) < R(4)) + 2
(48)
Comparing the first and third inequality in the above equations we get R(Φ) < + From the optimality of d* over the set (e) and since E 5 (e) it follows that R( I R(Q). Hence, R(0*)= < R(Φ) < R((*) + v. This completes the proof.
Next, prove Proposition 2 from the main body of the manuscript.
Proposition 8. For every 0, 0 and 6 E (0, 1), if Ha is finite hypothesis class, Assumption holds, K > 0, and if the number of samples [D] is greater than max 161." SL. } log ( (1
then with a probability at least 1 - 8, every solution & of EIRM (equation 6 is a approximation
of IRM, i.e. & E (c), R( < R( < R( + where tp" is solution to IRM (equation 5
Proof. From Lemma we know that if D is g-representative w.r. X, Hp. loss and distribution and if D is 5-representative X,, He, loss E and distribution P. then the the claim the above Proposition is true.
Define an event A: D is g-representative w.r.t X, Hg. loss E and distribution P.
Define an event B: δ is g-representative X, Ha. loss E' and distribution P..
161 8L } log (1HHel) R²
Define success as An B. Next, we show that if [D] is greater than max then P(An (B occurs with a probability least 5. P(AnB) P U > P(A3)
26

---

Published as a conference paper at ICLR 2021
P(B)). If we bound P(A*) I and P(B)) I 2. then know the probability of success is at least 1 5.
We write
P(4) ((4)=P() -)((p vh E Ho. ((()- - -1) 3h E Ho. (R.0) R(h)) > "}
P P{(D: D: 3h E Ho. [R(h) R(h)) 7) = (P U (D:I(()) R(h)) > -(())
51)
heHe
(49)
< P P((D: D: R(h) - R(h)) ((h)- 5
KEH
The loss function is bounded |((((-), -11 - L. From Hoeffding's inequality in Lemma it follows that
P((D ( Rh)) R(h)
5)
(-
DD|₂ 8
2 - 2 exp
(50)
Using this expression equation 50 in equation
49
we get
|D|2
8L2
4/Ho
2/H exp
I
D
log
(51)
8L2
2
y²
P(B) = (1B)=P((b)) Vh E H4. H4. ( (1)) - x)) = I-P(( 7h E Hg. (k) F >
))
P((D: 3h E Ho. Ho. (h) - H (h)) 3)
P U {D: - H -b)))()))A((b)>>>>> ")) he
(52)
Σ p ({D D: bEHo
(R: R (h) R (h) R(h)) 5))
The gradient of loss function is bounded (8(h(-)-))
Ow m w=l.ol S L'. From the definition of S 1 Recall that R (h) =
(h()),)) in equation 34 we can infer that / (l(().)) BlC (h, (X),,
From Hoeffding's 's inequality in Lemma] I it follows that
P(fD: R (h)
)) s 2xxp
(-
2 2exp (-
D|x2
(53)
81
167.4
Using the above equation 53 in equation 52 we get
| 16
5
16T.
4 ||/6|
Ho
exp
D
log
(54)
/
2
x2
Combining the two conditions equation 51
and equation 51 we get that if
DD max 161. 81
44 H44
log
2
then with probability at least 1 - event An B occurs.
27

---

Published as a conference paper at ICLR 2021
7.3.3 PROPERTY OF LEAST SQUARES OPTIMAL SOLUTIONS
We first remind ourselves of a simple property of least squares minimization. Consider the least squares minimization setting, where R(h)- = I[[( - h(.)
BI(Y~ - ((x))² E (Y 1 + || - xx
Ex B[((Y E |
- + Ex - h(x))²
(55)
Ex [Var[y ( [x]] + Ex
[F[v|x]- - xxx)))
In the above simplification, we use the law of total expectation. Both the terms in the above equations always greater than equal to zero. The first term does not depend on h, which implies the minimization can focus on second term only.
min R()) = Ex Var [x[]]
[
+ min Ex h
[ - x[
(56)
Assume that I' has full support over X. Define Vo E 2,1'(z) = E| = =x]. wh, R(h)
>
Ex Var [v]]] Since Re(h*) = Ex Vyr[y|xp x]] Therefore,
h" arg min Ex
(E [E[ A(X))*
(57)
Moreover, we conclude that h- is the unique minimizer. Observe that Ex [y[ - -(x)) is
zero for = h² From Theorem 1.6.6 in (Ash al., 2000), it follows that any other minimizer same as ht except over set of measure zero.
Recall the definition of m(x) = |E[
Lemma 4. Let I be the square loss. If Assumption 4 holds and m H4, then m uniquely solves expected risk minimization m arg mingeH. R(() and also uniquely solves IRM (equation 5).
- Eeec, x² (Q). From Assumption and the observation in equation 57 it fol- lows that the unique optimal solution to expected risk miniiization for each R. is Theretore, m also minimizes the weighted combination R.
Proof. R(Φ)=
To show the latter part of the Lemma, if we can show that E (e),,), then the rest of the proof follows from the previous part as we already showed m a minimizer among all the functions in Ha and SIV (c) C Ha.
Suppose m & (E). This implies there exists at least one environment for which | | R" (w- 22 > 0 + V=lw==1.l [w m) * O. As a result 3 w in the neighborhood of w = 1.0 where (" < (m) (if such a point does not exist and all the points in the neigbohood w = 1.0 greater than equal to (m) that would make V R% m) = O. which would be contradiction). Therefore, [w m) < (m). However, this is contradiction as know that m is the unique optimizer for each environment. Hence, m & StV(e) cannot be true and thus m (c). This completes the proof.
Next, we prove Proposition 4 from the main body of the manuscript.
Proposition 9. Let E be the square loss. For every v
0, > 0 and 8 (0, 1), if H4 is finite
hypothesis class, m E Ha, Assumptions 3 hold, and
if the number of samples [D is greater than max { (812 Dz
log( 4HH ), 16L
ot log(?)} then with
probability least 1 8, every solution & EIRM (equation 6 satisfies R(m) < R(Ø) R(m)+ + v. If also v < k. then Φt = m.
8I2 if the number of samples (D] is greater than " log( 2 [HHll every solution क to ERM satisfies R(m) R($*)
), then with probability at least 1 - 6, R(m) If also < RG then कt m.
28

---

Published as a conference paper at ICLR 2021
Proof. We first cover the second part of the Proposition. From Proposition know that the output of ERM will satisfy R( I R( - R( ))++))+>+) + v. In this case from Lemma 4 it follows that o+ From the definition of and the fact that < R implies that ot m. We now move to the first part of the Proposition.
For EIRM will derive tighter bound sample complexity than the one in Proposition 2 since we can now use the Assumption Observe that Ve E Err, V.l (w m) = 0 (see the proof of Lemma Therefore, R (m) 0.
Define an event A: D is "-representative w.r.t X, Ho. loss and distribution P. Define event B: D is such that (m) - R' (m)) S द. Since R' (m) 0, | (m) H (m)l s R (m)) < = R' (m) S 2. Therefore, m E (e)...
If AAnn B occurs, then R(m) S R()) S R(m) + is we justify claim next. Suppose & solves equation 6 If event A occurs, then from _ -rrepeesentative condition we know that R( S R())).
From optimality of & it follows that (()) < R(m) (event B = m E (e)). Moreover, from E-representative property, we conclude that R(m) S R(m) + z.2 We combine these conditions as follows.
V' R(क) - 7/2 R(( & R(m) S R(m) + V From the above we have R(m) & R(6) v R(m) + v. Recall the definition of i and since v < m.
(58)
Next, we bound the probability of success. B)AnB)= 1 1- P(A)) P P(B).
If can bound P(As) S and P(B9) s ş. then we know the probability of success is least 1 1.
We write
P(4) = 1000P((( P(D:: E Ho. R(h)- < ")1 = -1 1-P((P 3h e No. (R(A) (())-
")
From equation 5T if the condition
81.2
4|Ho|
D 2 1/
log
(59)
true, then true, then event A occurs with probability at most 4.
We write P -P(p(((()))))) 1, (m) R' (m)) - $
1) = P P((D::
( R (m)) > >m))>}}) $})
The gradient of loss function is bounded a)(())))) Lemma follows that
l S E.. From Hoeffding's inequality in
dwe
P(FD: (R R' R(()) > 3 exp(
ID|2
(60)
16I
We bound the above equation 60 by to get
DDEc2
8 D
16 2
2 exp(
log
{)
(61)
167.4
2
Combining the two conditions equation 59
and equation 61
8L2
16 v2 log
H44 )
}
D Z max
log
E2
This ensures P(An B) 2 1 - 8. This completes the proof.
29

---

Published as a conference paper at ICLR 2021
Before stating the proof of Proposition 5 will prove intermediate proposition. For clarity, will restate the result (Theorem from Arjovsky et al. (2019)) next.
Proposition 10. (Theorem 9 Arjovsky et al. 2019) If Assumptions and (with r = 1) hold and let & R (4h + 0), then
E
(62)
holds for all e E Eer ift t STay
Next, we propose an e-approximation of Proposition
Define eg = [EarT (w)min)) (12 - 2)
Proposition 11. Let be the square loss. If Assumptions 0<<< Eo. the solution &
(with 1) andd hhold, then for all
R Sc
(63)
satisfies 4 = (Ta), where E
Proof. Let start by simplifying V ow/w=1.0ft R ②), using square loss for e and linear represen-
tation 4 E R
dE
[Y w T xx))
w/ 1 4)
2E (( [ X( - 2E ( x
dw
w=10
= E x
2Φ E Y
(64)
Plug the above equation 64 in the condition R' (Φ) S to get
R' R'($) = Σ V w/ (w 3)||2 = Σ π + TT E
Xexe
(x
(65)
From the bound on T in Assumption 5 it follows that for each E Eer
E(Errl min
E [x क TT [x=ya]
(66)
E(Errl min
T E [x<xe -- TT ['
Note that if the condition above equation 66]is not true, then the preceeding condition in equation 65 cannot be true contribution from one term in the summation itself will exceed e. In the above equation 66 we are using the positive square since RHS has be greater than or equal to zero.
We compute the second derivative of loss w.r.t w
8 (Ye Φ x=w)² 21
a 2wEE [ xc)2]
2E T xeyee)
V Rw ④)
Ow
dw
(67)
= 2Φ T E X xrexe.1e
T
Since E" is symmetric, we can use the eigenvalue decomposition
= TAUT in equation 67
to get T n = & U To. Φ. Substitute & UT Φ to get Φb T = TT > Amin |
Amin ( t T T Amin 2 From Assumption 5 have Amin (E.) 2 Amin and from Assumption we have | 2 w. Therefore, we can deduce that 2 T > 2Amin Therefore, the second derivative defined in equation 67 always greater than equal to 2.
2 R (w ②) > 2\min
(68)
30

---

Published as a conference paper at ICLR 2021
Let é' = =2 / 4armin
/
We rewrite equation 66 in terms of V Vulu=11off ②) and é to get
V l v= 1.0
(w (w Φ)1 c
(69)
V 10 (w- ②) c
Since the second derivative (equation 68) is strictly positive and larger than or equal Amin
w in the neighborhood of w 1.0 at which V w/ w/w (w +) = 0. This holds for all the environments in E,, Define
c(w) V R (w 1)
Also, define (w) = ac(w)
Dw Also, è' (w) 2 RR (w- Suppose V Volw=100f (w-+)) = c(1 O.
Since for all w, e c(w)) 0 (from equation 68 3 w 1, where c(w²) = 0. Using fundamental theorem of calculus, we write
c(w) c(1) c[1)= ' c(u)duu 22 2mm 1)
Substituting w = w in the above
c(w") cc(1) 2 2w -
-c(1) 2 27min ( 1)
c(1) < 1 + 2)min
wº < 1
1+
(70)
2 2)min
2Amin
Suppose 1.. (w. c(1) 0. Using fundamental theorem of calculus write
d(1)-c(w) - c(w) = =f" c(u)du Aminwa(1 w)
Substituting w = wn in the above
c(1) c(w²) 2 22minw(1 -w c(1) 21 w²)
clEca gemi
c(1)
w
(71)
2Amin
2Amin
2Amin
E
Define 1/1
172
Combining equation 70 and equation 71 and
-18
1+ 2m
using the definition of mn and 7/2, conclude that D E 1 I+ If reparamterize
w T then 17 E [M2, m]. We expand this condition V w/ www 12 Re .)) = 0.
2wE i E [Yee T. (2)x
[ X
2E $ X T
Ow
T
E [x x
2w E² [x
TTT
Es [x
t: E [x T ST E. [X
(72)
24
E [x
(200 - (T>) [X
2 E
(m49) (T>) x) EE' =
T
E
1+ - STy) E [ = =0
1
31

---

1. From Published as a conference paper at ICLR 2021
Assume that 1 1+ty= # y, for all Err- If this assumption is not true and 4 1-17= - T for some Etr, then it already establishes the claim set to prove in this Proposition (since v7 E [12.7) ml)
( 1-1"" - -TT) E [x From the Assumption 6 and since dim(span{
Define oe =
Htp # ST 7, we know dim span qr) }} n
From rank-nulliiy theorem know dimension of kernel space of & (rank of ④ is 1) is equation 72 it follows that q" is in kernel space of Φ. Therefore, dim(Ker(O)) = n
dim span ge) ) S n which leads a contradiction.
Therefore, 4 1t1
y at least for one environment.
If Φ 7 + 7)), then | = >1* + 17722 > ||ST y1 + m2) N | 1 z w. In this simplification, we use (1 + 7/222 2
and ||ST>11 y112 > 2w (from Assumption 7
This ensures that for any solution of
1+ 1+V 1+V3-223 2/2
the form &b Ttm the assumption | N is automatically satisfied.
If = 71 then ||011 = ||ST
11 71? s ||ST>11
3-2/2 I Ω
the above simplification, we use (1 + m1)²
/elkeel
/ell This ensures that for any
2urA
2w
3-2v2 and y11 v 3+2V2"
Ω (from Assumption
1- 3-V3-222
solution of the form d 1497 the assumption || < Ω is automatically satisfied.
The entire proof so far has characterized the property of & that satisfies R (((b) S But how do we know such & exists. Recall Siy E Ho. For each environment Exr
who 1.0 R.0 (w 4) + TTT 11Ee(XXX" ( xx.ee110
E [X = -T [x - - SE [
=TE (1E2[Z99/00 2 1 - y
(73)
TET Z 2i 7 _
Zizz 1 T [
= 0
We use Assumption 5 in the simplification above in equation 73 Therefore, R' (Št)) = 0 and as
result the existence of a & that satisfies R' (Φ) < is guaranteed. This completes the proof.
first establish that Assumptions 5 and that Assumption 3 holds and we can thus use bounds and L' defined in Assumption 3
Before proving Proposition 5
are sufficient to ensure
Proving Assumption 3 for square loss l from Assumptions From Assumption
5 have
ze (7i, 72)
X SZe
(74)
| S ||||||
From Assumption 5 ||||| is bounded and | bounded - || is bounded as well (from equation 74). Therefore, there exists ssup 0c, X < ssup
32

---

Published as a conference paper at ICLR 2021
Y = 7) X t
| S 1|S 7|1|| +|
(75)
2
| V
xux + guup
3 + 2v2
In the last step of equation 75
|
2 Ω (Assumption 7) || S X derived 1+2
above (equation 74). and
ET' & E (Assumption 5). Therefore, Y is bounded and there exists a K such that y - K < Sup + suu,, Therefore, vD E He and for all Xe, y sampled from the model in Assumption we have
((9(X"),XX) = (Y" - x=xx'² <K + VDXx
(76)
al(w T x)
Φ x x- - Y) & (VII + K)
(77)
dw
w=100
From equation 76 we conclude that (()()),)) is bounded and there exists an L such that
[(((()).)) s L s (K VSXx='y?? From equation
77
conclude that ae(w a X,Y)
w=1.0
bounded and there exists an L' such that
E(wo (.).) Sw
L' s (VXX))))(VXX)) + K).
w=1.0
Define 6th 24-101 24-16y2 16
2 x
(win 2 and
3/Etz
Awin
2mmm
Next, we prove Proposition from the main body of the manuscript.
Proposition 12. Let be the square loss. For every e (0, eth) and ō E (0, 1), if Assumptions
(with = 1,.
v. hold and if the number of data points DD is greater than 161,
2² log x
(2Haol) then with
probability least 8, every solution & to EIRM (equation 6 satisfies & (STx) 7)a, where --ive---ve
Proof. Define an event A: {D: VT
VD E Ho. RR (Φ) R (+) s 2) If event happens, then
Á' (0) <
A(()) R (R)+R(())))
(78)
R'(0) < RŘ (+) R
3e A' (²()) 2
event happens, then every solution & to EIRM (equation 6) satisfies equation 78 From equa-
tion 78 that we should substitute with 3e and use Proposition 11 If the Assumptions 5
(with 1). hold and event AA happens, then for all
Eo the output of EIRM equa-
tion
6 Φ (ax), where a E [-
i-rvel Note that
14 1. Bemtie
Am Paal
l1tv
< 6g
= < Cth Hence, all that remains to be shown is that event occurs with a probability
at least 5. Next, show that if [D] happens.
>
161 log( 2[Ho]) ) then with probability 1 event A
Now need to bound the probability P(A). We will find an upper bound on the failure probability using Hoeffding's s inequality (Lemma] [ and the bound (derived in equation 77) as follows. We redo the same analysis was done in equation 54 for reader's 's convenience.
33

---

Published as a conference paper at ICLR 2021
P(A) = 1-P)(D-P((():) Ho.
(())))()))) 2 =P (P D j (0) (0) 2)
U
*EH
AiD P( PUU D () 2 7 Σ p P(D:: (0) - (0)) 1) s 2/H le
1827
$EH
ZiD 2 2/H le 1627
6 = P(A9) 5
[D] 2 log (2/Hao)
16L
P(AC) <
(79)
Hence, we know that if [D] 2 162 2
log (2 ), then with probability 1 - 5 event A happens.
The proof characterized the property of á that satisfies Ñ' (क) S But how do know such $ exists; we show that & always exists. Consider & that satisfiess the following.
R((D) S
R'(() R' (T) R ($) S
(80)
R' RT()) S c + + ( R (+)
Ñ (b) (follows from event A)
From equation 73 in the proof of Proposition 11 STyy Hg satisfies H' (Φ) <
we know that R (ST7) = 0. From equation 80
Since equation is a constrained optimization, a penalty based version (IRMv1) was proposed by Arjovsky et al. (211)) that minimizes R( + XR (+). Both Propositions 4 and Proposition 5 can be extended 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(VI/A). Define Am max(55² (5t2/3eth 1}
Corollary 1. For every 5 E (0. X > >th. if Asssmpttinns 6 (with r = 1). 7 hold, and [D] > max { 642
25 log A|Hal)) 32L"X² log(4) ( then with probability at least 5 every solution Ab
25
of IRMvI satisfies = (ST 7)a, where a E 1470
Proof. The empirical version IRMv1 minimizes
R( + x
Let us compute the risk achieved by the ideal invariant predictor ST%,
R(ST))= E (V=e
Σ ''Ee (y (
Σ
n (
T'E'
o² (81)
eEEt
e
Define event A: {D: (N(STT) R(ST 7) S $
event holds, then
R(STT) 7) S 0²²+ E
(82)
Define event B: {D:v0c Ho, ( [i < }
34

---

Published as a conference paper at ICLR 2021
If event B holds, then |R - RR((STT) 1 2 and if we plug in R (8T^) 0 (from equa- tion 73 then
n'' 75T 2
(83)
Define success as AnB. eventt AnB occurs, then from equation 82 true for any solution of IRMv1
and equation 83 the following
R( + XE' (o) S R(5 AR (ST 7) s 0²
c e
+
(84)
2
Since /() 2 0 it follows that
R' (0) - A 2/2 = [0) - 8² +
0² + 1
0² + 2 + e
2/4
R'
(0) S
(85)
In the last implication in the above equation 85
we use the condition that event B occurs. Let
a² x and substitute in the above equation 85 to get
R1( S
20
S
502
(since A 1,
0²
5
2
(86)
2X
2X
21
2
Recall Proposition 11 and see the condition in the Proposition is true, then every solution of IRMvl is (a)). where
3 2A
takes the role of E. If S
2X
S co
=
A
2
5e²
3ep
then
E
1 1-ToV
We arrived at the above result assuming A B occurs. We will now show that if [D] >
m 16 of log( 1 8L2X2
04 log(})), then with a probability at least 1 - 5, An An B occurs.
We write P(A) = P ((D D ()( R(ŠT))) S 2 1) - 1 P((D: (A(TTT)) - R(ST>)
$)) From Hoeffding's 's inequality in Lemma and it follows that
P (((D D: (R(ST)) - R()T)))> 2) 2exp (- ID|c2
8L2
2exp (- ID|2
= P(A9) - 8
8L2
2
2
8L2 2
4
> 5 2
|D] 2
log
P((A)
(87)
8
Next, we will show next that if [D] happens.
2
161
log( 1Ht ). then with probability at least 1 event B
2
Now we need to bound the probability P(B). We will find upper bound on the failure probability using Hoeffding's 's inequality 3nd as follows
P(B)=1-PP 1- ((D::((::) J0 Ho, (K) [0)) 1 =P -P(U U [D: (0)- R¹ (Φ) )
DEHo
P(
(4)
2 2 Σ (((D: R (0)) > ) s 2 le
165.
dEHo
EEHo
Z1D
8
8
2 174 le-
19
<
P(B9) S
2
2
14 /4
4 4
|D|
log
= P(
<
2
8
2
(88)
35

---

Published as a conference paper at ICLR 2021
Hence, know that if happens.
[D]
162 e²
log(
4 HHll), ), then with a probability at least
event B
Combining equation 87 and equation
88
we get if
16L
4 Ho
8L2 x2
(
4
[D]
z max
log
)
log
x2
then A n B occurs with at least
8 probability. Substitute
5 2X
and this completes the proof.
7.4 EXTENSIONS: POLYNOMIAL MODELS, INFINITE HYPOTHESIS CLASSES, CLASSICCLASSIIFCATIIN
7.4.1 POLYNOMIAL MODEL
In the main body of the work, we did not cover the polynomial model in detail due space limita- tions. In this section, we study the polynomial model and show how the results for the linear model can be generalized this case. We first begin by stating the model itself. Assumption 8.
e (ategorical(T")..)
Ove Esr
y = (z)) + t5,,5" 1 Zi, B[ca] O, (,E(( = 1 X = S(Z;, Z
(89)
Assume that Zi component of S is invertible, i.e. 35 such that Š(S(Z7, k2)) = Di and also Sty # O. the above Çp is polynomial feature map of degree p defined as cp R2 + R" where denotes the dimension of the input the map Sp. cp (W) = [W, W W W...., (W & W...p times &
W) 1( p=z1] and @ is the Kronecker product. Also, a = Ei a' Ve Etr, TT 2 |Ear1'
The support of distribution of Z =Z7 Zg}, Pz.. is bounded and the operator norm of S, ||||| (ax (S) (A maxxx (S) is maximum singular value of S), is also bounded.
Define 2e = (7i, 75) and say Zi E IRC and 7 E TRd.
Can directly the analysis from the linear case? We cannot directly 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 ( /==
where S is a block diagonal matrix with diagonal matrices defining it
given {s, 58,
S0p
Define X= 53 (x where 5 (X*) is the polynomial feature map of degree of dimensional input X Similarly, define Zi (5(Zi) Zf). where (5(Z)) the polynomial feature map of degree p of с dimensional input Zf and define Z Ço <ctd(2 (2*), where (2*)(2") is the polynomial feature
map of degree of c - d dimensional input Z (2 2z). Observe that each component of Ž is also in Z.
From the model we know that X = 55 We would like to remind the reader of the mixed product property of tensors. Consider matrices A E [iix B E kxi C E xx D gxx
(A (C & D) = (AC) @ (BD)
(90)
In the expressions that follow, we exploit the mixed- sixedd-roodcc property of Kronecker product stated above.
SZ. - diag (s
1= = 013 |=1 = [xxx3)) = (X = X (91)
-
36

---

Published as a conference paper at ICLR 2021
Define 3 = diag ($01)
SX 552 diag
(5 /=
diag
(58))
(Z0.01)p 1=11
diag (($5)
(Z
1==1 ((552))00)
[(2)))0)) -
(5(2)) Zī
(92)
The dimensionality of X is ni' Ei Assumption 9. Inductve bias. Ht is finite set of linear models (bounded) parametrized by o RT', 5T E Hg. w > 0, > 0, V E H., I |111 S and 2w S >12 S 342/2 S
2p11_ n-1
We next compute the norm of S in terms of the norm of S. Recall are using operator norm defined as ||S|| = max (S). [|S]] = diag (5 y
Since $ is diagonal ma- |||||* (Laub, 2005). Therefore,
trix ||||| max;- I....p [ ||}. Also, note S
||S|| = max 1...... p {| Hence, if ||S|] is bounded, |||| is also bounded.
| Observe that | = Z || Hence, if || is bounded, | is also bounded. Since X 5Z% We can conclude that the same line of reasoning in equation 74] equation 75 equation 76 equation 77 to conclude that the loss and the gradient of the loss are bounded.
Also,
X is also bounded. We can now follow
We rewrite the above model in Assumption 8 as linear model in terms of the transformed features.
Categorical n ) x" OVe E Y 7/ + o. ee 1 71 E[{²] 0, [({() = 52,1 sup
(93)
X S2
We showed above in equation 92 that Z defined in equation 91 component of Z is invertible,
552. Zi. We have also shown above that support of Ze is bounded and the norm of $ is bounded and as a result || and the loss and the gradient of the loss (conditions in Assumption 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. set of training environments Ey said lie in linear general position of degree for some & > n - n /r and for non-zero I E R and
dim span
E
E
X ET
CEE. > n -
(94)
We denote =[X'
Assumption 11. For all the environments e E Exr, Ee is positive definite.
Define the minimum eigenvalue over all the matrices E. as Amin minege,, X(S.e). Define Esh 2 min
Berl (w:Amin)².
From the analysis in this section, we see that we have been able to construct a linear model identical where the role of Xe, n, S, S, ΣS, Amin- éth is taken by X 2., S, 5e, Xmin- é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.
5
Proposition 13. Let E be the square loss. For every e E (0, 8 E 1). if Assumptions 10 (with r 1), 17 hold and if the number of data points D is greater than 161.
- log Re
then with a probability at least - 8, every solution & to EIRM (equation 6) satisfies = 18T where E (1 v-/vel
37

---

Published as a conference paper at ICLR 2021
7.4.2 INFINITE HYPOTHESIS CLASSES
the work so far we have assumed that the hypothesis class Hg is finite. this section, we discuss infinite hypothesis class extensions. Before do that state important result covering numbers that we will use soon.
Lemma 5. (Shalev-Shward & Ben-Davie 2014) Define set A = {a RR,, ||al| I A Covering number for r-cover of A given as N, (A) is bounded as
(₂(A) < A2
(95)
A. Infinite Hypothesis Class: Confounders and Anti-causal variables
In this section, we seek extend Proposition 5 to infinite hypothesis classes.
We restate the Assumption 7 for linear models.
Assumption 12. Inductve bias. Ho is set of linear models (bounded) parametrized by & E R
H4 { R" 0 < | < 12. TT E Hy. 2w S //5 3-2v2 Ω
Note that the only difference between Assumption 7 and Assumption not required to be finite anymore.
12 is that the hypothesis class
We already established in equation 75 equation 76 equation 77 that from Assumptions 5 and 7 can show that the conditions Assumption Shold, Mioodd,ee, i.e., loss and the gradient of the loss are bounded, and also X y bounded. The same conclusion follows from Assumptions 5 and Assumption 12 Hence, for the rest of this section, we can state that | S Xuuu, Y - K, the square loss
is bounded by L'. the next lemma, we aim to show that if Assumption and 12 hold, then R' (4) Lipschitz continuous.
(()()))) :) is bounded by and (((w)()))).))
dw
1=110
Lemma 6. If Assuupttion 5 and 12 hold, then R' (Φ) is Lipschitz continuous in +
Proof. The output of the model | <
| S "ux (From Cauchy-Schwarz).
R' R'(0) Σ B Φ
x= yo)
R ( - R((02) = Σ
x E
1 ( yro)² -B" o (xx
- ( -roc
TTx - ra] -- ( [Jxx -xol
regl
E
(96)
We bound each term in the summation in the equation 96 above
38

---

Published as a conference paper at ICLR 2021
$ x ( x - yo]]² -vo] -E xx -xxe---ell
Re
[0]x== - r9] - E. 1 (0) - - -x))) (0x0x
-E.
E
12 x (0 - xrr))
1 11 -x - r9]
EE
- E $2 x ( x |0]
+ K)
19 x - r9)
E.
$3 [0) x - ro] + F. x0 vr)
Ea
4 2/x' - r3]
2 2V (v + KK)
- ( x (0) - - -. + E -0|X"=xxx''' x5x*) 22 + K)
E
E
(0|x-- oo xx - yr) 2 </ + K)+
E.
2 (VSXx K)
(97)
We bound each term in the last line in equation 97
E
- 23xx -
r
$ S [ - क.. xxxx) (0) - rr
(98)
s (V + K)E" | - 22) Xx
1 (V + K)E [| 22 x2]|X
(Cauchy-Sewaad
$ (VNX= + x) 22
E- 1 X ( X - [x
o[x* E 1 X [x]x==
VΩX
(0) 42) x0|
Cauchy-Seware
(99)
"!X""
[| x
VS( b| T21
Substituting equation 98 and equation 99 in equation 97 get
R (..) R' ( < 2 ( (22VXX K) T21
(100)
Therefore, R' is Lipschitz with a constant C' < <2V"(x""))VVXX" + () + 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 E E (0,eth thh) and d E (0, 1). if Assumptions (with r = 1), 12 hold and if the number of samples [D] is greater than 321,4 log ( }3)] then with a probability at least - every solution d to EIRM (equation satisfies (T). where littve: T-vel
/i6/ Vin
c n log
39

---

Published as a conference paper at ICLR 2021
Proof. Following the proof of Proposition 5 goal is to compute the probability of event A:
(+) S 5)- We construct minimum cover of size (Ho)) (See Lemma (VP Ho, R ( R
5) with points C {
Compute the probability of failure at one point 4 in the cover
P[[D. IR R ( 2)
AiD 2e 32274
(101)
We use union bound to bound the probability of failure over the cover C as follows
p { max
R (b,
R (+)))
1 < 2N,
32274
(102)
SCEC
Now consider any o E Ho and suppose Φ is nearest point to it in the cover.
$ { [()-
( R (Ø,) R (2,)) 2n'C'
(103)
In the above simplificati,, we exploited the Lipschitz.ommiiuitty of R'. Therefore, for each E Hap
|R (+) - ( - max 8CEC
R (1, - R (2.) + 2m7
(104)
[ Rl (+)
max R (9) R' (2,) 21/
max EEHo
3C
2 m
Set
sc in equation 104 and from equation 102 with probability at least 1 - NN,(H))2e 321
max EHH
R (1) n' (Φ) S 2/2 (since max 0,EC
[R' (Φ,) - R' (Φ,) S E/4)
(105)
We bound N, (H4)2c
33674
and solve for bound on D to get
8L
22,, (No)
DD|
log
(Use Lemma 5)
8
(106)
32L/4 /4 [D] >
16C'
VIn
2 )
n log
+ log
2
Therefore, if condition equation 106 holds, then event occurs and following the same argument as in the proof of Proposition 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 con- tinuous functions that we formally define next. Define map + : P x X' - R from the parameter space P and the feature space X' to reals. Each p E P is a possible choice for the representation (p.-). Consider neural networks as example, P represents the set of the values the weights of the network can take.
Assumption 13. $ P R Lipschitz continuous function (with Lipschitz constant say Q).
such that Vp P, ||a|2 < P. X' C IR is closed and bounded, thus there exists a Xsup < so such that Vo E X, |||||| s X
C I is closed and bounded, thus there exists
Y C R is closed and bounded, thus there exists K co such that Vy E,, |v| $ K
40

---

Published as a conference paper at ICLR 2021
Assumption 14. Lipschitz loss and gradient of loss. R(D) is Lipschitz with constant C, R (Φ) is Lipchitz with constant C'
From Assumption derive the conditions in Assumption and Assumption 14
Px.X IR
is 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 x))² & (M+ Hence, there exists L such that ((((()),)) < L & (M + KYP
86(w ().))
is bounded:
Belw +.).)
IY ((p. (p) X)1 & (K + M)M.
8vc
w=1.0
Ohc
w=1.0
Hence, there exists an L' such that
a (-).) th
S L' S (K + M)M
w=10
Lemma 7. If Assumption 13 holds, then R(v(p,)) :) and R' (1(p, -)) are Lipschitz continuous in p.
is Lipschitz:
[R(((P.)- ((((g.)))) Σ (() -x)))))))
E
rejii)
T"
E
(D X²) - - - -E
-
(0(9. - y-
E
- 2/ - x9) ES 1
X.) (c)
+ 2(g, X(. 2Y
E'
@(p, 4(9. x
E. (p, + 1(2. x²) 2Y
2(M -K) (9) - 2 |
T'E
2112(M + KQQ
(107)
Therefore, R is Lipschitz with a constant C < 2(M + K)Q R' is Lipschitz:
Σ क [ (p, (9)))) - yo] -E00 r3)
|R (2(p,) R (2(g )) =
n E (0), - r- E XX - -
(108)
We bound each term in the summation in the equation 96 above
41

---

Published as a conference paper at ICLR 2021
-.oe)-ro))) (F,,))/)00),,00000000000 =1 +
-Ei -
E
p(p. x
E
(g, X5) (,,,
E
Φ (2),,
|m] -x)
E X0) (0(.. - rr)l) 2M + K))
E
P(p. (06,, r-)
E X).. x9) - rd] + + x,) -v
(1 (4) 2M KK)
(p(D, X²) - 1/1 x (
- -x--))))) + (x, - (l. 2M + KK)
(0(p, x.) 1 ((g.)9) -210) 22) K)+
E
E.
X0) x9 1u xx)x 2M (M
(109)
We bound each term in the last line in equation 109
- ((),)xx)))-((,xx)(()))x.))x)) - xY)]
s (1,1)/""""---,,,,00000,,,,,,1 - Y
(110)
$ (M 1-)0K)Ppp--- 41
-
$ X)) 2 X
(1II)
$ MQ - 4l
Substituting equation 110 and equation
equation 109 get
IR (o(p,)) - R (Q ))1 < 2M + + all
(112)
Therefore, R' is Lipschitz with a constant C' s 2M + (2M + K)Q.
C. EIRM: Sample complexity with no distributional assumptions
In this section, we discuss the extension of Proposition too the infinite hypothesis class case. Con- sider the problem in equation and replace the with k. Define distance between the two sets S (c) and its approximation ST (c+ R) as follows
dis(k)
max gES"(ett) (EtK)
min d(g, h)
(113)
where d(g, ht) is some metric that measures the distance between functions and h. Observe that <. dis(a) dis(b).
Assumption 15. lima- dis(x) 0
Define
322L2 D2
16CVPk
+ log (2)]
8L"4 x²
SC'
+ log (2
(2)1}
D = max
k log
k log
V
R
Proposition 15. For every > O and 5 E (0, 1), if Assumption
13 15 hold, then E > 0 such
that if the number of samples [D] is greater than D* then with a probability at least 8, every solution d to EIRM (replace e with e + x inequation 6) in SIV + 2x) and R(Φ) R( ) $ v, where th" is solution IRM in equation
42

---

Published as a conference paper at ICLR 2021
Proof. We divide the proof in two parts.
Define event A: {D: Vp E P, P, (+(p.))) R (p(p, s k).
In the first half, will show that if event A occurs, then 8tV (c) C S x)) - S + 2n) and then bound the probability of A not occuring.
(o(p.))) Sc = ((())))) - ( + Sc =
(114)
( Set ((+ ( = (9)
Therefore, (c) (c)
( 1)) (o(p,))) SCtk = - + R' (D(p, )) S + + (2(p.)) (9 = (0(((p, )) 2x
Sct
(115)
Therefore, 8tV (E + x) S (e + 2x)
bound the probability of event A not occurring. Using the covering number (from Lemma 5 construct minimum cover of size (P) with points C1 {p, 1-
we
Compute the probability of failure at one point py in the cover
P
[[ D (0(p,, )) A (0(ps, )){>>>>>>>
K. ]
2e
8277
(116)
We use union bound to bound the probability of the failure over the entire cover C₁ as
p[[ max pyEC1
Rt (o(p). :) - f (0(p) )) ] < <N,(P)2e
627
(117)
Now consider any E P and suppose Pj is nearest point to it in the cover.
(+(p,))))) )) (0(pp :))
RR' (o(p,)) 1) R' (o(p).)) R (.(p, R' (#(p,))) + R' (F(p,,)) R' (*(,,)) (118)
< R (D(p).)) - fi (0(p,,)) :)) 2//C'
In the above simplficiation, we exploit the Lipschitz continuity of R'. Therefore
vp E R (o(p.)) R (0(p, -)) max R (o(p), )) N (0(p), 1) 2:jC' E,
(119)
R (o(p. R' ( )) < max IR (I(p,, R' (T(p,, 1) + 2//C max PEP
P,EC1
2
Set 17
x
in equation 119 and from equation 117 with probability least 1 N2
82/4
4C"
max
R (d(p.))) --' ($(p, < (since max pEp
[R (i(p,, -) R (o(p,, 5) x/2) (120)
P,EC,
We bound N,(P)2e
and solve for bound on DD to get
8T4
(NN,(P))
D >
log
R²
(121)
8 14
8C' VPk Pk
D
>
klog
+ log
R2
43

---

Published as a conference paper at ICLR 2021
Therefore, if condition in equation 121 holds, then event A occurs. If event A occurs, then (c) _
(I x) C S (e + 2n).
(()",-) is a solution to IRM equation and it satisfies (p E (e) R({()")) < (p,
Define an event B: {D vp E P, (p -) R( >> - 5}
If event occurs, then for a solution ( -) of equation (where is replaced with x) satisfies
R(* )) V' S (( s ((v(p",)) R(P(p", .-)+ R(*(p.)) - ())((p")))))))))
(122)
Using the covering number (from Lemma we construct a minimum cover size N2,() with points C₁ = {-
Let us bound the probability of failure one point Pj in the cover
[[(((())) - - (1)0(, 4
2e
- 3262
(123)
We use union bound to bound the probability defined as
P
{{
D: max [R((())-)) R(((p, 1)) 4 -
"
20D (P)22e
326
(124)
PJECi
Now consider any E P and suppose pj is nearest point to it in the cover.
[R(((p,)) (((((p,))) ))]
[)(((p))) R(*(p), 1) R(*(p),)) R( +)+ R( R(((p,)))
(125)
< <R(()) - iR(2))) +
Therefore, for each p P
R( :) (( -)) S max [R((()) )) R(D(p; +2n P,EC1
(126)
max [((((p))))) :) R )) < max R(((p,,)) R( )) + 2m/C pEP
.2
Set 17 = 80
y
in equation 124 and from equation 126 with probability at least (P)2ee
38774
max (p(((p.)) A' (4(p. ))) 5/2 (since max pEP
R (P(p,) :) R (o(p,, vv44 (127)
PjECi
We bound (2)()2e 320.
and solve for bound on |D| to get
32I.2
2.N, (N,
D >
log
2/2
8
(128)
32
16CVPk
+ logg 2 $
D
>
& log
Therefore, if condition in equation
128 holds, then with probability at least
event A occurs.
Also, if
32I,2 n²
16CVPk
2
8L x²
8C' 8C'V
+log (2) (129)
> max
k klog
+log
)
I: klog
44

---

Published as a conference paper at ICLR 2021
then the event A n B occurs with least 1 probability.
Project (p,-)) (E + k) 8t (e), i.c., find the closest function in terms of the metric dis, to
obtain ((ū .). If event A occurs, then p ( (e 2n). The distance |/ ill S dis(2x).
[R( )) R( )) - C أ1 S Cdis(2x))
(130)
We choose Kiy such that < Ku (use Assumption 15 such that Cdis(2x) S
R( ( - R( +
(131)
Therefore, by combining equation 122 amd equation 131 we can conclude that if event n B (pp ))] S v. From the conditions on
occurs, then ( :) S 2x) and R )) R
[D], we know that A B occurs with probability 1 6. We substitute ((p",) ;)) as 4 and ( as i and this completes the proof.
D. OOD Performance: Covariate shift case
In this section, we discuss the extension of Proposition the infinite hypothesis class case.
Define Dj
324
k log
32CV/kk + llog (t)] 162 log (3)
Define D
321 23
k log
[16C/Pk]
+ log }
Proposition 16. If Assumptions 4 13 hold, m. Hy and if the number of samples \D] is greater
than Di, then with a probability at least 8, every solution & to EIRM satisfies R(m) s R(Φ) < R(m) +
If Assumptions 13 hold, m E H.p and if the number of samples [D] is greater than Dz. then with probability at least 1 every solution o² of ERM satisfies R(m)= I R(D*) I R(m) + v.
Proof. We begin with the first part. Following the proof of Proposition 5 our goal is to compute the probability of event A: {Vy E P, [R( R( )) < 2) Using the covering number (from Lemma 5) we construct minimum cover of size b = (P) with points C = {p, )
Compute the probability of failure one point py the cover
P {{D [R(0(P)))) - M1
2²
< 2e
(132)
We use union bound to bound the probability of failure over the cover C
p[[D: DD: max [R(0(25) )) (((((ps)))) D1 < N2(P)2e
321
(133)
Now consider any p E and suppose pj is nearest point to in the cover.
R[((p.-)) R(((p.)) =
[(((()))) R( + R(O(p).)) R( R( (((((,,))
(134)
$ R( R( 1)1 2n/C'
In the above simplification, we used the Lipschitz continuity of R. Therefore
Vp E R( )) R( ))1 S max RR :) R(((())))))) 27CC
EyEC
(135)
max (((((p))) A( )) < max [R(((p).)) R( }) 21C PEP
P,EC
45

---

Published as a conference paper at ICLR 2021
Set sc in equation 135 and from equation 133
with probability at least 1 - (P) 321/4
It (L(p,-)) R (0(p, )) - v/2 (since
Ii' (D(p), -R (1(p),)) m/44 (136)
max pEp
PrEC
D 32F.7
We bound (2,2)
< and solve for bound on [D] to get
81.2
Nn((P)
[D] 2
log
(Use Lemma 5
(137)
321.2
16C
2
D 2 k log
+ log
Therefore, if condition in equation 137 holds, then event A occurs. Observe that the optimal solution p" for expected risk minimization p" E arg min R))(p.)) satisfies (pp"..) m(-) (From Lemma 4 and E H4). If event A occurs, then from same argument used in equation 48
solution p of ERM satisfies R( R(0(p", )) S R( :)) R($(p", .)) + true.
We now move to the second part. From the first part of the proof, we conclude that when
32L2 log 16C + log ]
DD 2
(138)
with probability least
event A occurs.
Define an event B: is such that R (m) - R' (m)) $ 1. Since R' (m) - 0, (R' (m) R' (m) s
R (m) S 2 R' (m) S 2. Therefore, E (c)
We write P(B) PPLD: j (m) - H' (m)l S + = P(D /R (m) H (m)) > }) The
gradient of loss function is bounded a)(+()))
1==10
$ T.. From Hoeffding's inequality in Lemma
iit follows that
P(I: (1)) (1))
< 2 exp (
DDE2
2
16
DDEE
2 exx
S
16T.4
16 D
log
(139)
Combining the two conditions equation
138 and equation 139
=222" 22
+32C Pk
]
4
16 2 log
4 }
IDI 2 max
l klog
+ log
This ensures P(An B) > 1 8. If event A n B occurs, then follow the same justification as in the proof of Proposition 4 to claim that solution ş E P to EIRM equation 6 satisfies 2((((",.)) S R +) < (("(p". +
This completes the proof.
7.4.3 EXTENSIONS BINARY CLASSIFICATION (CROSS-ENTROPE
In the main body of the manuscript, we focused on regression (square-loss). In this section, we discuss the results that 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 an order that makes for easier exposition for the classification case.
46

---

Published as a conference paper at ICLR 2021
We begin by showing how extend Proposition 4 to binary classification (cross-entropy). Recall
that the entropy of distribution Px is (P)) = - -E log(dP)] Recall that the cross entropy of
Q relative to IP is (I, Q) = -Ep log [og(
H(P) + KL(dPP] The cross entropy loss e "X) is the probability of label conditional on X²) is given as ((((x)) = log (J(X)) + (1 Y+) log (1 - ((x)))
for binary classification when using a predictor f - [0.
For the discussion below Q(Y=|x*) is defined terms of f as follows Q( 1|x*) f(X*) (Q(Y= = x)) = ((X*).
(f) EE [((Y,, ([(Y,,((') [xx))
Ee ye ye log (J(X*) (1 1og (1-
(1- (1-
TR. E. [= log (f(X*) + (1 [[[ log (1 - [X-](
EE [P(Y* = (Xx)) (f(x)) + P(Y = 0|xx) log (1 (1- (1-)(x))
(140)
= E (=(P)11119))0()11X
+
= E [[(P((*X))) + |
From the above it is clear that Q(Y91*)= = -P(Y)X**)) minimizes the risk in an individual environ- ment.
Assumption 16. Invariance w.r.t all the features. For all e, E Eaul and for all X, E[ x] = Y = z). X ~ xx and Ve Eall support of Fxx is equal to X.
Observe that in the clasffccaassiiicattinn setting the above assumption amounts to equating the ditional probabilities ( (X*)) and [x))
Recall that map m (from equation 2 simplifies Vo E X
m(x) = E [Y* = P 1|x =
(141)
If Assumption 16 holds, then from cross-entropy decomposition in equation 140 it is clear that m solves the OOD problem (as is optimal w.r.t each environment). It is also the unique minimizer. We can justify it based the same argument presented in equation 56 Suppose there 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 positive thus contradicting the optimality. This shows m is the unique optimizer. The rest of the arguments presented in the proof of Proposition earry over this case. Therefore, Proposition extendss to the cross-entropy loss.
another
Note that Proposition 2|s proof was agnostic 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] [o, 1]. We could not generalize Proposition 5 to cross-entropy loss and that left as future work.
Next, we move to showing how Proposition I can be generalized to binary classification.
Assumption 17. Existence of an invariant representation. 3 xx X - Z such that Ve, E Eall and Vx E X, (=))) = E|Y (*)((((()) (x))
Recall m defined in equation
Yz E x² (X)
ml 2
E [
Z = == = 1|X = =)
(142)
we get the following. Define composite predictor w o क* Substituting f = o 5' in equation 140
For the discussion below, a distribution -( is defined terms of w t as follows R(Y==
1 1 w (X*) (R(YY= 0|X*)= = x0 (X*)).
(wwo = E. ( [( + E (X(RX
(143)
47

---

Published as a conference paper at ICLR 2021
If all the data is transformed by @", then from the above decomposition equation 143 is clear that 2( = P(Y+IZ") the optimal predictor for each environment. Hence, w P(Y
1 1 is the best choice for
Assumption 18. Existence of an environment where the invariant representation is sufficient. an environment e Eall such that Y 1. z", where Z =X.....
We derive relationship follows for the environment q satisfying Assumption 18
P(Y (²) = P ( (follows from conditional independence in Assumption 18
(144)
Also, note that since Z Φ5' (Xx) we have
2" = P(Y
(145)
From equation 144 and equation 145
we have
P((Y" = P
(146)
We use equation 146 the cross entropy decomposition from equation 140
(f) = E" H (P(Y")))) E (L(PP(")
(147)
Recall Q(Y" = ) = f(x) (Q(Y" = 0) = 1 f(x)). Also, recall w² (29) P(Y 22). From the above clear that f = = क* is the optimal predictor for environment q.
(9 o -b) = E [(((9)))))
(148)
The expected conditional entropy for environment is defined H E (P (2)) is the
risk achieved by w² Φ. Also, H 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 g. which satisfies Assumption [8
Assumption 19, Ve Eall, He S 11sup, A19
Hus
Therefore,
Ro (w" o x0) = suu
(149)
From equation 143 for all the environments
(" +04) H
(150)
Observe that max.6E.. Rd (w= o xa) = Hsup, From the above assumption it is clear that for all predictors f 10,10
Vf. max R 2 R* SuP eECall
(151)
min max Rt > Sup eEEall
Since max R" (5") Haup we conclude that t5 is the predictor that solves the OOD problem in equation This completes the extension of Proposition to cross-entropy.
7.4.4 ON THE BIASEDNESS OF ERM
Consider the model in Assumption 5 For each environment e E &tr, define a vector p [ Define matrix á with p as column vectors p = ip',...,, Define vector -= [r] T/ E.rll] where recall from Assumption 5 x is probability point comes from environ-
ment o.
Proposition 17. If Assumpttio hollds Ho linear hypothesis class with parameter To and if the rank of p at least one, then ERM is asymptotically biased, i.e., even with infinite data ERM will not achieve the desired solution STa, except over a set of measure zero of probability distributions
48

---

Published as a conference paper at ICLR 2021
Proof. Consider the case when ERM has access to infinite data, i.e., solving the expected risk minimization problem stated as mingen. R(Φ). We will consider the linear model in Assumption
5 and assume H4 linear hypothesis class parametrized by & R". We simplify the V R(() for
the square loss below
VoR(Φ)
Σ eEErr
x (ye ((re
T. xxxx)
(152)
We compute the gradient for + ST as
01o==ST
R($)
Σ
T"
(Use Assumption 5)
eEEt
(153)
Σ <
x E'E² E x
Recall P¹ = E _- X and p Ip'...... pl Recall म - [ ....,... 1
Setting the gradient
defined in equation I53 to and using the above matrix notation we get
0,1 1,
(154)
If स satisfies equation 154 then ERM is unbiased, else it is not. Consider the set of vectors in the probability simplex { T
1, क 2 0} and define a uniform probability distribution over it. of the columns of is non-zero. As result uniform random
Since rank of f 0 at least draw from this set of probablity distributions would have zero probability of satisfying pḌ = 0. Therefore, STy is not the optimal solution to ERM and thus the solution of ERM would be biased away from ST
In Proposition 5] we had assumed Assumptions 5 6 hold. If we also assume that rank of is at least one, the Proposition continues to hold. then the rank of ý is at least one. From proof of Theorem 10 in Arjovsky et al. (219), linear general position continues one of the E X
least one e E Evr, E XX" is non-zero,
hold except over set of covariance matrices with measure zero even when is non-zero.
Also, in the above Proposition 17 we only required that rank of p least one. However, make the additional assumptions Therefore, if Assumption 5 hold and rank of p is least one, then ERM asymptotically biased and IRM can be within /E neighborhood of the ideal solution with the sample complexity shown in Proposition
the result of the above Proposition continues to hold.
49


Reviews:
[{"review_id": "rW91QK0qBDt", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\n\nThe authors provide sample complexity study of the Invariant Risk Minimization (IRM) in the case of scalar predictions and one dimensional representations. Instead of considering the penalized version of IRMv1 proposed by **Arjovsky et al, 2019**, the authors study its constrained counterpart introducing a relaxation into the requirement of zero gradient. The main conclusion that I understood from the work is: the performance of IRM (provided the assumptions are satisfied) will never be much worse than that of ERM and will be superior to that of ERM in some particular cases. The authors support their theoretical discoveries empirically.\n\nReview:\n\nThe paper is rather dense and the exposition could be improved. I had some troubles in Section 3.3.2 since the spaces of the variables and matrices are not introduced (but it *might* be standard in causality literature, which I am not familiar with).\nIt is difficult to judge the significance of the work. On the one hand, the proofs of most of the results follow fairly standard tools of empirical process theory and I am not sure how much more insights this work provides compared to that of Arjovsky et al, 2019. On the other hand,  the statements themselves seem to be novel and they fit well in the big picture.\n\n\nPros: The theoretical results on IRM are novel and the authors introduce new variants of colored MNIST.\nCons:  The proofs follow the same pattern as that for ERM. The fact that IRM is constrained does not bring much novelty in the proof. See for instance: **Woodworth et al, Learning Non-Discriminatory Predictors, 2017** and **Agarwal et al, A Reductions Approach to Fair Classification, 2018** for similar proof techniques. It is also unfortunate that the main body refers too often to supplementary material for some details that in my opinion are very relevant to the main body. For instance, empirical version of the gradient, additional discussion on why ERM is biased in Section 3.3.2, and the setup of the experiments are all postponed to appendix.\n\nQuestions:\n1. On feasibility: I wonder whether there is an issue of feasibility for the problem in Eq. (5) and Eq. (6)? Is it obvious that if the set $\\mathcal{H}_{\\Phi}$ is finite then we can always find a feasible representation?\n2. Slack variables: what is the reason to introduce slack variable $\\epsilon$ in the formulation of Eq. (5)? It seems to me that it is sufficient to introduce it only for Eq. (6), since the authors only need to show that $\\Phi^*$ is feasible (with high probability) for Eq. (6) in order to prove their upper bound. I also do not see why $\\kappa$, introduced on p.5, cannot be zero.\n\nSuggestion: \nThe phrase **''``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''** is not convincing. Finite class is not that realistic in practice (correct me if I am wrong). If all the results of this work can be stated for infinite classes of bounded capacity (e.g. VC-class or totally bounded spaces), I suggest to completely erase the case of finite classes. \n\n\nAfter rebuttal: I thank the authors for comprehensive rebuttal, which addressed most of my concerns. I update the score accordingly. \n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "bParcyRequ2", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "SUMMARY\n#######\n\nThe present paper proposes a sample complexity study to compare Empirical Risk Minimization (ERM) and Invariant Risk Minimization (IRM).\n\nIRM is a framework developed by Arjovsky et al. 2019 to avoid learning a predictor based on spurious correlations (e.g. color background of a photo).\n\nData are assumed to be sampled from different environments, and it is assumed that there exists a transformation $\\Phi^*$ of the input data $X$ such that the expected value (and the variance) of the output $Y$ given $\\Phi^*(X)$ is the same among all environments. If $\\Phi^*$ is the identity, this assumption encompasses covariate shift.\n\nUnder the additional assumption that $\\Phi^*(X)$ is sufficient on at least one environment, authors show that for the square loss $x \\mapsto \\mathbb{E}[Y \\mid \\Phi^*(X) = \\Phi^*(x)]$ is a solution to the Out-of-Distribution (OOD) generalization problem (prop. 1), that aims at minimizing the maximum environment risk.\n\nIRM is obtained by restricting the search space to invariant predictors, i.e. predictors that write as $w \\circ \\Phi$ such that $w$ minimizes over all $u$ the environment risk of $u \\circ \\Phi$ for all environments. This condition is further replaced with a gradient constraint. Empirical IRM (EIRM) is achieved by replacing expectations with empirical averages.\n\nAuthors first show the concentration of the solution to EIRM around that of IRM (prop. 2).\n\nThen, under additional assumptions, authors study how EIRM approximates the OOD solution by exhibiting sample complexities under both covariate shift (prop. 4) and more involved distributional shifts (prop. 5). Results suggest the superiority of EIRM over ERM in the later setting.\n\nExperiments on variants of the colored MNIST dataset are presented, matching the theoretical findings. Code is provided.\n\n\n\nCOMMENTS\n########\n\nThe paper is globally well written and understandable (despite notation getting heavier and heavier with the assumptions). IRM is clearly introduced and the motivations well exposed.\n\nThe subject tackled is of great interest, and the theoretical angle adopted is novel and insightful. Experiments are eloquent.\n\nAbout Asm. 1, how hard is it to check in practice? Does someone have to explicitly exhibit $\\Phi^*$? Maybe an example where $\\Phi^* \\ne \\mathsf{I}$ could be useful at this point.\n\nAbout eq. (1) and (3), the maximum over all environments has been replaced by the sum. This can be seen as replacing a $\\| \\cdot \\|_\\infty$ with a $\\| \\cdot \\|_1$, which has a very different behavior (tends to cancel components, while the infinite norm makes all components small). How can this change be justified and impact proofs? The same question could apply to $R'(\\Phi)$ that is the sum instead of the max (but the $\\epsilon$-relaxation makes this change more natural).\n\nAbout eq. (3), how to be sure that invariant predictors exist?\n\nIs the restriction to real-valued $\\Phi$ necessary? I feel it makes the compositional architecture $m \\circ \\Phi$ a bit obsolete. It rather appears now as a trick to use the differential constraint (the $\\nabla$ could be replaced by standard derivative since $w$ is one-dimensional). Also, is there a way to consider sub-differentiable, or non-differentiable loss functions here?\n\nAbout Prop. 2, can authors elaborate a bit on the $\\kappa$ constant, that seems to relate to $\\epsilon$ in a complex way, and is likely to be very small?\n\nAbout Prop. 4, as far as I understood, one of the main changes is that $\\epsilon$ now replaces $\\kappa$ (that I think is better to get an idea of dependences). Otherwise, the bound seems almost identical to that of Prop. 2, suggesting that the difference between EIRM and OOD solutions is \"the same\" as that between EIRM and IRM solutions?\n\n\n\nMINOR COMMENTS\n##############\n\np.3 Asm. 1, $\\exists$ in the beginning of an english (i.e. non mathematic) sentence is a bit weird\nP.3 we define *the/a* invariant map\np.4 Prop. 1, If $\\ell$ is *the* square loss\np.4 for *the* cross-entropy loss\np.4 the parentheses two lines after seem unnecessary\np.4 the space*s* of representation and\np.4 eq. (4) *s.t.* omitted\np.5 the statement of Prop. 3 is a bit misleading, putting the reference in the title would be clearer (although authors clearly cite the reference)\np.6 Asm. 7, induct*i*ve bias and & symbol quite unusual\n\n\n\nOVERALL EVALUATION\n##################\n\nThis paper proposes a very interesting theoretical study to elucidate which one of ERM or IRM is better and when. Authors also present conclusive experiments matching the theoretical findings. For all these reasons I think this contribution deserves acceptance.\n\n\n--- EDIT POST REBUTTAL ---\n\nI thank the authors for their detailed answer(s) and the efforts they put in editing the revision (in particular I agree with other reviewers that Prop. 17 adds clarity). Overall my stance on the paper has not changed and I still recommend acceptance.", "extra_scores": {}, "extra_text": {}}, {"review_id": "RWH5vyYgShM", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: 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.\n\nGood points:\n- focus on an important problem, and answering elements for recent debate in the scientific community e.g. https://arxiv.org/abs/2010.05761 [scientific importance]\n- general insights clear and easy to grasp (Table 1) [clarity of results]\n- remarkable effort to explain technical results (e.g. paragraph \"implications of Proposition 4\") [readability effort]\n\nQuestions:\n- it seems section 3.3.2 and generally results on the confounder and anti-causal scenarios are restricted to polynomial mechanisms ($g$) ?\n- in experiments with the confounder case CF-MNIST (Fig. 1b) it seesm IRM only achieves chance-level performance (=~ .45 test error) ? \n\nPoints limiting the relevance of the paper:\n- overly technical paper: it is hard to follow for readers interested in general (or deep) ML but not literate in generalization theory\n- density of the paper: most of the space is dedicated to technical results which albeit interesting are seemingly \"shoehorned\" into 8 pages. I wonder if the conference format makes justice to such a deeply technical work... 47 pages total with appendix\n- difficult to check the proofs: all are in appendix (no sketch) but there are so much of them. As a reviewer I find it discouraging as the time needed to check correctness would be much more than what I can afford in a short cycle for a conference such as ICLR\n- result a little over-sold in abstract/conclusion: it seems that the theoretical results are somewhat limited (see question on sec. 3.3.2) and that is omitted in both abstract and conclusion, maybe giving a false sense of generality to the results. Maybe the studied case - which is already impressive - is more general than I think but I'd like a comment of the authors about it.\n\nPoints that would improve the paper (and my score):\n- clarify the meaning of limitation of result in section 3.3.2 + add it in abstract/conclusion\n- comment the CF-MNIST (Fig. 1b) experiment\n- add/rewrite if possible \"interpretation for the layman\" paragraphs (e.g. last paragraph of sec. 3)\n\nFinally, I'd like the authors to think about the best venue to make justice to their work. It seems to me that a general ML journal would be a good option that would 1) allow more space to make the text more self-contained 2) allow more time for proof checking and discussion/optimization 3) allow more interpretative comments to be included in the work and make it accessible to a wider audience. In the mean time a preprint could also circulate to share the current research results.\n\nalso: small typo on title of Assumption 7", "extra_scores": {}, "extra_text": {}}, {"review_id": "HFkHqaBU6MX", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Summary\n\nThis 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.\n\nThe main object of study is empirical invariant risk minimization (EIRM), which minimizes the empirical risk under the further constraint that the predictor is (approximately) optimal in each environment. The main result is that ERM and EIRM have the same sample complexity when the environments are related via a covariate shift, but that under other situations (namely confounding variables and anti-causal relationships), EIRM continues to enjoy guarantees, while informal arguments are offered as to why ERM doesn’t. Experiments are given to illustrate this disparity, and thus potentially justify the use of EIRM.\n\n### Strengths\n+ The paper provides a very methodical and detailed reasoning to try and identify the benefits of invariant risk minimization and provides several compelling arguments.\n+ The main strength of the paper is in placing the finger on the situation where ERM and IRM appear to differ. It doesn’t complete the story here, but at least identifies the reasons behind prior accounts that seemed to conflict on the surface.\n+ The math is laid out very clearly and the gradual introduction of detail from the main text to the supplements keeps the presentation of the ideas manageable.\n+ The new colored MNIST datasets proposed in this paper is a very useful contribution, which more realistically captures the kinds of shifts and dependences that may arise in realistic ML applications. \n\n### Weaknesses\n- In the quest for detail, some of the high-level perspective is lost. The first place where this happens is when the paper claims that Assumption 1 is equivalent to covariate shift when $\\Phi$ is the identity map. But the constant-variance requirement of A1 does not follow from covariate shift, so in fact A1 is a stronger assumption in that case.  And this is not mild: it is the glue that binds together the proof of Proposition 1. (Incidentally, the argument for the latter can be given more simply as: in the sufficient environment, the optimal estimator enjoys the constant variance as MSE; if the objective is higher then, in all other environments, we can reduce it by reverting to the conditional expectation; thus the latter must be the optimal predictor.)\n- The second place where the high-level perspective is lost is when tying together the ideal OOD predictor (1) and the IRM formulation (3). There are two facts that distinguish (1) and (3). First, (1) does not require the predictor to be optimal in each environment unlike (3). In fact, in the manner of most minimax results, the solution might sacrifice even in a “good” environment by being worse than optimal, so that some of the “bad” environments fare better. Second, (3)’s average objective may result in a choice of $\\Phi$ far from the ideal representation. To see this, imagine two environments that need a more elaborate representation and a third that doesn’t but is very noisy. (1) would favor the latter and give a simple representation that would perform well in the worst case, while (3) may ignore the third environment and focus on the first two and thus optimize the average case.\n- This blurred high-level perspective makes it so that the rest of the paper analyzes the procedure in (3) (or rather its gradient-constrained formulation) while striving to recover the solution of (1) (as given by Proposition 2). It’s a conceptual disconnected which unfortunately only becomes evident upon reading through, and is not adequately highlighted by the paper. It’s also possible that the bulk of additional assumptions introduced, as well as the loss and model structures adopted, are really a result of enforcing the reconciliation between these two different paths.\n- The entire methodology of EIRM hinges on a perfect awareness of the environment (the constraint cannot be calculated without it). In that sense, it is quite an unfair comparison to plain ERM. In fact, the awareness can be thought of as adding a particular inductive bias to ERM, so it is not much of a surprise that the analysis presented is not more novel than regularized ERM. \n- ~~What the paper claims a lot but does *not* deliver is a lower bound for ERM under the cases where new upper bounds for EIRM are given. In Table 1 there is a “No” under OOD, on page 7 we read “[t]his shows that EIRM works in more settings than ERM”, and in the conclusion we read “we proved that […] ERM can be asymptotically biased”. But apart from informal arguments and experiments, this is not in fact shown.~~ In particular, considering the great number of assumptions leading to Proposition 5, it is not at all clear how ERM behaves under the very same assumptions (note that the experiments do not satisfy all these assumptions, e.g. $\\pi^\\min=0$, so we cannot treat them as counterexamples).\n\n### Comments\n* The $1/\\kappa^2$ dependence on $\\kappa$ is a bit weird. $\\kappa$ is the smallest $|R’(\\Phi)-\\epsilon|$ over all $\\Phi$. So in principle there maybe a $\\Phi$ that makes $R’(\\Phi)=\\epsilon$ and Proposition 2 vacuous. I think the way it is used is the same as the reason behind why $\\kappa<\\epsilon$: it’s because we assume a feasible solution of the optimization exists, so $0 < R’(\\hat \\Phi) < \\epsilon$, thus $|R’(\\hat\\Phi)-\\epsilon|<\\epsilon$. But $\\kappa$ is *not* defined as $|R’(\\hat\\Phi)-\\epsilon|$, perhaps it should be?\n* After Assumption 5, the paper should make it clear that the initial focus is on $S$ being linear.\n* Some things should be moved up from the supplements for completeness. For example, the definition of your constraint, the $\\theta$ for Figure 1d. \n* Typos: Assumption 7, *Inductive* bias. $[\\psi^1,\\psi^e,\\psi^3]$: change $e$ to $2$. Figure 1, CMNIST for all.\n\n### Overall\nThe paper considers interest settings of distributional change and corresponding learning formulations. It makes a lot of assumptions to obtain sample complexities that justify the use of empirical invariant risk minimization, and falls a bit short by not giving a formal converse for the inadequacy of plan empirical risk minimization, despite making the claim. Nevertheless, the contributions are insightful, and the paper may be worth sharing with the community.\n\n_[Edit: I'm glad the author clarified some of the tensions between the different definitions. The ERM bias proof seems also okay, but it's still not clear how ERM performs with all the additional assumptions of the IRM result. I'm still not on the same page regarding $\\kappa$, but I'm willing to take a leap of faith on the bits that I'm uneasy about. I would suggest this paper is shared with the community, since I believe it has enough good insights.]_\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper considers learning settings with distributional change. It makes a lot of assumptions to obtain sample complexities that justify the use of empirical invariant risk minimization, and falls a bit short by not giving a formal converse for the inadequacy of plan empirical risk minimization, despite making the claim. Nevertheless, the contributions are insightful, and the paper may be worth sharing with the community.
The grading were overall positive from the reviewers, though particularly critical, and I doubt the whole paper could be fully double-checked: one could question the ability of the reviewers to perform a deep analysis on a 48-pages theoretical paper in the time constraints imposed by a conference model... 
Decision: Accept (Poster)