Out-of-Distribution Optimality of Invariant Risk Minimization
Authors: Shoji Toyota
Kenji Fukumizu
Venue: Accepted by TMLR ()

Abstract: Deep Neural Networks often inherit spurious correlations embedded in training data and hence may fail to generalize to unseen domains, which have different distributions from the domain to provide training data. M. Arjovsky et al. (2019) introduced the concept out-of-distribution (o.o.d.) risk, which is the maximum risk among all domains, and formulated the issue caused by spurious correlations as a minimization problem of the o.o.d. risk. Invariant Risk Minimization (IRM) is considered to be a promising approach to minimize the o.o.d. risk: IRM estimates a minimum of the o.o.d. risk by solving a bi-level optimization problem.  While IRM has attracted considerable attention with empirical success, it comes with few theoretical guarantees. Especially, a solid theoretical guarantee that the bi-level optimization problem gives the minimum of the o.o.d. risk has not yet been established.  Aiming at providing a theoretical justification for IRM, this paper rigorously proves that a solution to the bi-level optimization problem minimizes the o.o.d. risk under certain conditions. The result also provides sufficient conditions on distributions providing training data and on a dimension of feature space for the bi-leveled optimization problem to minimize the o.o.d. risk.

Published in Transactions on Machine Learning Research (01/2024)
Out-of-Distribution Optimality of Invariant Risk Minimiza- tion
Shoji Toyota
shojt
The Institute of Statistical Mathematics
Kenji Fukumizu
fukumizu
The Institute of Statistical Mathematics
Reviewed on OpenReview: Attps: // openreview. net/ met/ forum? id- pWs JDa
Abstract
Deep Neural Networks often inherit spurious correlations embedded in training data and hence may fail to generalize to unseen domains, which have different distributions from the domain to provide training data. Arrovsky et al. (2019) introduced the concept out- distribution (o.o.d.) risk, which is the maximum risk among all domains, and formulated the issue caused by spurious correlations
minimization problem of the o.o.d. risk. In- variant Risk Minimization (TRM) is considered be a promising approach to minimize the o.o.d. risk: IRM estimates minimum of the o.o.d. risk by solving bi-level optimization problem. While IRM has attracted considerable attention with empirical success, it comes with few theoretical guarantees. Especially, a solid theoretical guarantee that the bi-level optimization problem gives the minimum the o.o.d. risk has not yet been established. Aiming at providing a theoretical justification for IRM, this paper rigorously proves that a solution to the bi-level optimization problem minimizes the o.o.d. risk under certain condi- tions. The result also provides sufficient conditions distributions providing training data. and a dimension of feature space for the bi-leveled optimization problem to minimize the o.o.d. risk.
1 Introduction
Training data used in supervised learning may contain features that are spuriously correlated to the response variables of data. Deep Neural Networks (DNNs) often learn such spurious correlations embedded in the data and hence may fail to predict desirable response variables of test data generated by a distribution that is different from the one to provide training data. To list a few examples, in a classification of animal images, models obtained by conventional procedures tend misclassify cows on sandy beaches because most training pictures are captured in green pastures and DNNs inherit context information in training (Beery et al. 2018 Shane 2018). Another example is learning from medical data. Systems trained with data collected in one hospital do not generalize well to other hospitals; DNNs unintentionally extract environmental factors
specific to particular hospital in training (AlBadawy et. al. 2018 Perone et al. 2019 Heaven 2020).
Arjovsky et al. (2019) introduced the concept dut-of-distribbuiin (o.o.d.) risk to formulate the issue caused by spurious correlations. Let X' and y be mcasurable spaces of explanatory and response variables respectively. Let & be a set with each element e E called the domain (or environment) e. Assume that for given domain e E E, there corresponds a corresponding random variable (Xe, yε) that takes values X' y with its probability law Px.,ye. Assune we are given training datasets D {(xf, vf Pxeye i.idd from multiple domains Err C. E. For given predictor X +
R²()== (dd

---

Published in Transactions on Machine Learning Research (01/2024)
denotes the risk of f domain e. The o.o.d. risk of the predictor f is as follows:
Ro.dd (f) = max R² R*)f).
(1)
cEE
formulated the problem caused by spurious correlations as a minimization problem of the o.o.d. risk U
which is the worst-case risk over & including unseen domain & - Etr- Arjovsky et al. (2019)
min R...d. (f): JEF
(2)
where F is the set of all measurable functions f + y
It is difficult to directly solve the o.o.d. risk minimization 2 since we can not evaluate the maximum of risks among all domains &, including unseen domains & Etr, only by data from training domains &,r C E. Invariant Risk Minimization (IRM) is a rapidly developing approach to the challenging o.o.d. risk minimiza- tion (Arjovsky et all, 2019). Its proposed predictor f := w T is composed of two maps: feature map Φ : X + H. which is called an invariance, and predictor w H + + which estimates the response variable of feature ((x). Here, for a given feature space H. we call a measurable function 4 ; X + H. an invariance when it holds that (x = Py (XX=2) for any e,, 62 E { maps by solving the bi-leveled optimization problem
Arjovsky et al. (2019) estimated the two
mingeZI.wEw Σ (wo))
eEEer
where W is a model of predictors w H - y and Itr is the set of invariances captured by training domains Eyr:
I, = { X + H. Pye = (x for any E1, 61, E E
(4)
Influenced by the seminal study, several alternative bi-leveled optimization problems have been proposed
(Ahuja et al. 2020
Chang al. 2020 Ahuja al.
2021a b
Lin et al. 2022a Zhou et al.) 2022 Liu et al.
2021a
Lu al.
2022 Koyama & Yamaguchi 202
Parascandolo al.
2022
Krueger
202
Toyota
& Fukumizu
2022
Lin al. 2022b Huh & Baidya
2022
Rame et al. 2022
Pogodin et al. 2023 Chen et al.
2023
Tan et al.
2023 For example, Ahuja et al.
2020
proposed a novel bi-leveled optimization problem
leveraging the principles of game theory. The recently proposed Maximal Invariant Predictor [Koyama &
Yamaguchi 2021 employed a new bi-leveled problem grounded in the concept of information theory.
While IRM is widely recognized as promising approach for the o.o.d. risk minimization 2 it comes with few thcoretical guarantees; especially, a mathematical guarantee that the bi-level optimization problem (3) gives the minimum of the o.o.d. risk (1) has not yet been established [2] The original IRM paper did not mention any theoretical properties for the minimum of (3). Rosenfeld et al. (2021) proved that, assuming that data follow a simple linear Gaussian structural equation model (SEM), a predictor obtained by (3) makes prediction relying only on a feature of Xe E X' whose distribution does not depend on domains Rosenfeld et al. 2021 Section 5). However, their analysis did not focus on relations between the bi-level optimization problem (3) and the o.o.d. risk (2). More recently, Kamath et al. (2021) provided an example of distributions on which a minimum of (3) does not minimize the o.o.d risk (Kamath et al.) 2021 Section 4). However, their analysis assumeed that data follow particular SEMs constructed to derive the case where does not provide minimum of the o.o.d. risk; for verifying the o.o.d. performance of the bi-leveled optimization problem (3) it should be analyzed under more general assumptions distributions.
Aiming at providing a theoretical justification for IRM, this paper rigorously proves that a solution to the bi-leveled optimization problem also minimizes the o.o.d. risk I formally speaking, we prove that the
The definition is based on conditional independence Peters et. al.
2016 Koyama & Yamaguchi 2021 Kojas-Carulla et al.
2018). while Arjovsky et al.] (2019): Ahuuj et al. (2020) used a different type of invariances basedl on arg min,,. R (wo instead
of Py (x Throughout the paper, argue by adopting the definition based conditional independence.
SSuce it difficult solve the bi-leveled optimization problem
8 several papers have propused optimization methods for
such IRMv1 (Arjovsky ct al. 2019) Invariant Rationalization (Chang et al. 2020 While their optimization ability
solving should also be discussed theoretically, this paper does address it. and only locuses on the problem of whether,
assuming that can solved completely, the resulting predictor minimises the o.n.d. risk.
2

---

Published in Transactions on Machine Learning Research (01/2024)
inclusion
arg min Bew.
Σ eEEt
R. (w o b) C arg min Ro.od (f)
(5)
fer
is attained under certain conditions. The result also provides sufficient conditions on the training domains Etr and the feature space H. to minimize the o.o.d risk. In our analysis, we set distributions on domains & by the ones proposed in (Rojas-Carulla structures, unlike existing theoretical analysis of IRM (Rosenfeld et al. 2021 Kamath et al. 2021). and they
al. (2018) The distributions do not rely on any specific SEM
are used for the analysis of methods related to invariances Carjul-Carulla et al. 2018 Toyota & Fukumizu 2022
The rest of the paper is organized as follows. Section 2 illustrates two main theorems. Section 2.1 provides
the first main theorem, which states that the inclusion (5) achieved in the regression case. In Section 2.2 we extend the first theorem to the classification case. The novelty and significance of these two theorems
are discussed in Section 2.3 We provide review of the prior works concerning the relationship between the bi-leveled optimization problem
3 and the o.o.d. risk [ in Section 3 The two main theorems stated in
Section 2 are proved in Section
Section is devoted to brief concluding remarks.
2 Main Results
We explain the settings and assumptions persisting Uhroughout our analysis.
We set domains (Xx),"))eee by the ones proposed in Rojas-Carulaa
Carulla et al. (2018). Let ; Xi x No where X1 := Rd and X₂ IR with d2 E Noo: and (X(,Y') be fixed random variable on X, x y. Rojas-
Carulla al. (2018) defined the domain set by all the probability distributions with the fixed conditional
distribution Pysp namely, denoting T.x, X' + X, a projection onto X2, {(Xe, Vo)eeEe is defined by
(X*,)") YEEE (X,))) random variable on X x x (x) Pr
(6)
Note that, under the setting 16) the projection tx X + Xi is an invariance among &, because Prje Pr for any YX E ((X)Y))))e For simplicity of theoretical analysis, we assume that the conditional distribution Prijx! has a probability density function p (1)2z))
We explain assumptions about the feature space and models. The feature space H for invariance $ E Ier is assumed to be the multi-dinensional Euclidean space Rdx, Moreover, assume that & E Ier and E W in the minimization problem (3) run only continuous functions; namely, we investigate the property of a
solution for
min
Σ (w o +), eEEr
(7)
SET We
where WCc is the set of all continuous functions w H y and Iee is the set of continuous invariances captured by training domain Err:
_ee = ( X X H| Preijo(Xe1) Pyez ((x=) for any el, e2 ey, continuous
2.1 Case I: Least Square Loss
First, we consider the case where y Rdy (dy E Nyo) and I is the least square loss; that is, for a given predictor f + -,, its risk Rf) on (Xe, ((xx((())ee..... Yo)eee given by
RT := f ( fl00///
The following theorem ensures that the optimization problem minimization problem 22) under four conditions:
7 provides solution for the o.o.d. risk
Theorem (o.o.d. optimality of the bi-leveled optimization problem under least square loss setting).
Domains ((Xe, YeEe are assumed to be 0 We also assume that the following four conditions hold:
3

---

Published in Transactions on Machine Learning Research (01/2024)
(i) = where Ico is the set of continuous invoriances captured by all domains &, not training domains Esr:
Io Io= ( : - H.| (x = (x for any e1. 61. E E, & continuous }.
(1i) UnEe, supp (ax (xo)) Xi. Here, for probability meusure y on X1, supp(u) is defined by supp(u) == (ar E Ri [No, open neighborhood around T1 = ((NN,)) 0}.
(10i) The dimensions di and dx on the subspace X1 C X of the impul space X and the feature space H = Rdu satisfy d1 S dh.
(tv) Pr has continuous probobility density function (y((z)). Here, we call p' (y((21) continuous when correspondence X1 x y E (21.9) + (y)(x1) is continuous.
Then, we have
arg min SET weWco
Σ
Rw b) C arg min Rood (f).
(8)
fEF
eEEtr
Here, F is the set of all measurable functions f f : + y.
We explain the feasibilitics and interpretations of the above four conditions.
Condition (i): Condition (i) implics that invariances captured by training domains Eer correspond to the ones by all domains E. AArjvvsy et al. (2019) also discussed the relationship between the equation I,, = I and o.o.d. generalization, briefly illustrating that the equation Itr I facilitates the estimation of predictor with high o.o.d. performance solely based on data from training domains Eir (Arjovsky et al. 2019 Section 4.1). If it holds that Z, II we can capture an invariance & E I among all domains & only using the training domains Err. |Arjovsky et al. (2019) pointed that, once an invariance & E I among all domains is oblained, a predictor w" that minimizes risks only on training domains Err, namely w" E arg minl Eeee.. R o Φ); satisfies R (w= = min,, Rº (w o +) for all domains e E &, including unseen domains E- EEr,, under certain settings. Developing the discussion by Arjovsky al. (2019), Theorem [ clarifies a more rigorous relation among the equation I, I, the o.o.d. risk (1) and the bi-leveled optimization problem 77 the equation I = Iy is one of the sufficient conditions for the bi-leveled optimization problem to minimize the o.o.d. risk 1.
The condition I, = I is not generally satisfied and Peters et al. (0016) Arjovsky et al. (2019) presented sufficient conditions on the training domains Err for the equation I, = when data follow simple SEMs. Peters et al. (0016) proved the equation Itr = I holds when distributions on domains follow a linear Gaussian SEM and training data are obtained by certain types of interventions (Peters et al. 2016 Section 4.3). Arjovsky et al. (2019) generalized the result by [Peters et al. (2016). Assuming that data follow lincar SEM, which is not restricted to a Gaussian distribution and certain type of interventions, Arjovsky et al.] (2019) deduced a sufficient condition for the equality I, = I on training domains E,r, which is called lying in the general position (Arjovsky et al. 2019. Assumption 8). On the other hand, sufficient conditions for the oquality I, = I under the sctting (6) have not yet been revealed. Providing them would be an important area for future research.
Conditions (ii) and (iii): As shown in Lemma the conditional expectation f yy-p' (y/zz) E|Y u = 21] achieves the minimization of the o.o.d. risk, signifying that the information embedded in Xi is important for predicting response variables training domains Eer covers X1 that contains such important information for o.o.d. prediction. Condition (ili) implies that H is such a large feature space that a feature $ : X + H can preserve information on the X-ccomponent of æ e X by selecting & appropriately. Condition (iii) also provides a practical perspective on how construct the feature space H Rdw; the dimension dx on the feature space H should be fixed high. The dimension du of the feature space is fixed by hand, and hence, Condition (iii) is expected hold unless fix the dimension of the feature space too small.
unseen domains. Condition (ii) implies that the support of
4

---

Published in Transactions on Machine Learning Research (01/2024)
Condition (iv): Condition (iv) presents contimuity of the p.d.f. of Pyr[x By Condition (iv), also have continuity of the conditional expectation y' assume that the model W consists continuous functions, and hence, Condition (iv) ensures that the model includes the conditional expectation [X]), which minimizes the o.o.d. risk (Lemma 3
('(lai)) = E = y|X = z1). our analysis, we
2.2 Case II: Cross Entropy Loss
Theorem can be easily extended to the classification case where W has probabilistic output and evaluate risks by the cross entropy loss. Let y be a finile set (1,..., m) (m E N>o), and we model w H + y by H - Py, where Py denotes the set of probabilities on }; namely
py P P
Σ Ps 1
Here, R4 = {x E 2 2 }} and denotes the i-th component of p. We call H. - Py continuous, that is pg E W when correspondence H. + (hh E R is continuous, seeing ph(h) E Py as a vector on RD', For a given po H + Py and E V, (po(h)); often abbreviated by po pi(h)). The risk evaluated by the cross-entropy loss is then written
Rs (po o 4)= =)=}- f
We expand Theorem 1 to the above classification case:
Theorem (o.o.d. optimality of the bi-leveled optimization problem under cross-entropy loss setting).
Domains {(X", }eEe are assumed to be o Assume that, in addition to (i) (iii) in Theorem 4 the following condition (v) holds:
(v) For any of E X1, # {u E (1 > 0 >.
Then, we have the inclusion
arg min $EI PeE
Σ
R o b) C arg min R Roo.d (f),
(9)
P
eEEtr
where F is the set of all measurable functions f X Py
Condition (v) indicates that domains {(Xe, Y'EE have high uncortainty in labels y y given xı E X. The condition is expected to be feasible when classes y are subdivided and difficult to be uniquely determined from 21 E X1
2.3 Novelty and Significance of Theorems and
Theorems and and their proofs have the following four novel and significant points:
Setting of Domains The first point is the setting of domains. The setting by Rojas-Carulla al. (2018), which is used throughout our analysis, does not impose any specific SEM structures, linearity, and Gaussianity on domains while existing works on theoretical analysis of IRM assumed that data follow simple SEMs. Theorems 1 and indicate that, under such a general setting, IRM presents the minimum of the o.o.d. risk. This implies that our results provide solid foundation to use IRM for broad range of o.o.d. generalization problem.
3Tee as definition of continuous, we call f F measurable when correspondence X + f(x) R! is measurable, seeing (xx é Py
RV
5

---

Published in Transactions on Machine Learning Research (01/2024)
As well as the assumption domains, prior theoretical results of IRM assume that the underlying true distribution Pyzpx; is represented by a simple SEM (Arjovsky et al., 2019 Rosenfeld et al. 2021 Kamath et al. 2021p. On the other hand, Condition (iv) only imposes Py'|x; the contimuity; hence, Condition (iv) is significantly mild condition in comparison with the assumptions on Pyi|x] by prior works.
Assumption the Underlying Distribution Pynxx
Characterization of Invariance Second, a theoretical characterization of invariances 4- 7 is given in Lemmas and 6 it is proved that $* I can be represented as $- V T for some continuous map V.. Any theoretical characterizations have not yet been presented, and hence, the results in Lemmas [ and 0 are novel. To present the non-trivial characterization, the proof by contradiction. Lemina 4 and 6 play an important role in our desirable assertion and
develop novel theoretical technique based
hence, the derivation of these lemmas is a significant technical contribution of our analysis.
Range of Invariance The fourth point is a range of invariances Φ: we assume that & run all continuous functions, while most of the existing works on theoretical analysis of IRM assume that tb run more simplified functions, such as linear functions (Rosenfeld et al. 2021 or variable selections (Toyota Fukumizu, 2022) It is common to construct a learming model of invariances with deep neural networks in the context of IRM, and hence, the variable selection and linear function settings by Toyota & Fukumizu (2022). ; Rosenfeld et al.
(2021) p are significantly simplified to analyze IRM. On the other hand, our large class of continuous functions is relatively realistic compared to existing ones, since it is widely recognized that neural networks of sufficient
size can represent a wide range of functions (Cybenko 1989 Hornik et al.
1989
Barron 1993 Mhaskar
1996 Sonoda & Murata 2017).
3 Previous Works
As explained in Section Rosenfeld et al. (202) Kamath et al. 2021 derived the theoretical results
concerning the minimum of the bi-leveled optimization problem (3 and its connection to the o.o.d. risk [
Rosenfeld et al.] (2021) proved that a predictor obtained by minimizing 3 predicts Y E y relying only on
a feature of X X' whose distribution does not depend on domains Rosenfeld et al. 2021 Section 5). However, they did not provide any connections between the minimum of (3) and the o.o.d risk. Morcover,
they assume that data follow a linear Gaussian SEM, and that invariances in the bi-leveled optimization problem run lincar functions for simplicity. Unlike their analysis, this paper derives the direct relations assume that data follow the distributions proposed by Rojas- Carulla et al. (2018) that do not rely on any specific SEM structures and that invariances run all continuous functions including neural networks. Kamath et al. (2021) provided an example of distributions on which a minimum of (3) docs not minimize the o.o.d risk. However, the distributions are particular SEMs constructed to derive the case where 55 is violated, and analysis in more general settings is required (Kamath et al., 2021] Section 4). In construct, the distributions by [oooRRooaraaa Carulla et al. (2018) used in this paper do not rely on any specific SEM structures, and they are used to analyze estimation methods related
between (3) and the o.o.d. risk ( Additionally,
to invariances (Rojas-Carulle et al. 2018 Toyota & Fukumizu 2022).
Arjovsky et al. (0019 Koyama& & Yamaguchi (2021): Rojas-Carulae Carulla et al. 2018 discussed theoretical re- lations between invariances and the o.o.d. risk 1b. As explained in the last section, Arjovsky et al. (2019) intuitively explained that the condition I,, = I facilitates an estimation of a predictor which can predict Y
on unseen domains only by data from training domains &rr (Arjovsky et al.] 2019. Section 4.1). They also derived sufficient conditions on training domains for the equation Ier I, assuming that data follow simple linear SEM (Arjovsky et al. 2019 Theorem 9). Koyama & Yamaguchi (202I); Rojas-Carully et al. (2018) presented sufficient conditions for invariance & to achieve the minimum of [.. Koyama & Yamaguchi (2021) proved that the invariance that maximizes the mutual information with labels also maximizos the
o.o.d. risk. Rojas-Carulla et. al. (2018) proved that, under the domain setting (6), the conditional expectation
| (X*) = _1] also minimizes the o.o.d. risk, even when E| (Xx) is nonlinear. However, all the results by Koyama & Yamaguchi (2021); RRojas-Carulln et al. (2018): Arjovsky et al. (2019) did not deal with any theoretical connections between invariances obtained by minimizing the bi-leveled optimization problem
and the o.o.d. risk (1). Tt does not follow obviously that the minimum of 33) satisfies these sufficient.
6

---

Published in Transactions on Machine Learning Research (01/2024)
conditions by Koyama & Yamaguchi (2021) Rojas- Caaas-a et al. (2018). and hence our main theorems not be deduced as a trivial corollary of the results by Koyama & Yamaguchi (2021) Carullaa et al. (2018). To discuss the non-trivial relation between the bi-leveled optimization problem 3 and the o.o.d. risk I we cstablish a novel characterization of invariances (Lemmas and and derive the main theorems
based on it.
To reduce the annotation cost required for the original IRM approach, Toyota & Fukumizu (2022) introduced
new bi-level optimization problem similar to (3) They considered situation in which the training data for target classification are provided in only one domain, while the task of a higher label hierarchy, which requires lower annotation cost, has data from multiple domains. Under the availability of data, they deduced bi-level optimization problem, in which invariances were given by additional data in a higher label hierarchy. For further details, we refer the reader to the original paper [Toyota & Fukumizu (2022). Their study provided a detailed theoretical analysis concerning their method and its connection to the o.o.d. risk; however, they did not analyze relationships between their bi-leveled optimization problem and the o.o.d. risk, which is the focus of this paper. Instead, they investigated relationships between optimization method for their bi-level optimization problem and the o.o.d. risk. Moreover, they assume that invariances & run all variable selections for simplicity of theoretical analysis. On the other hand, this paper derives the direct relations between the minimum of the bi-leveled optimization problem and the o.o.d. risk ( Moreover, we consider the more realistic setting for the analysis of TRM where invariances & run all contimuous functions.
4 Proofs
In this section, we prove Theorem and 2. Through the section, for X E X' and a E X, its Xt-components (i = 1,2) are denoted by X; and z; respectively.
4.1 Proof Sketch of Main Theorems
Before giving rigorous proof, we briefly describe the rough proof sketch of the main theorems. The following two lemmas (A) and (B) play an important role in our proof:
(A) The conditional expectation X[
E x and conditional probability Py- Py
minimize the o.o.d. risk under the least-square and cross-entropy losses respectively (Lemmas 3 and
(B) Any invariance among all domains can be represented by the composition of the projection onto X;; that is, E To can be represented as
o t
for some continuous map v (Lemmas and 6
The two Jemmas intuitively conclude the proof of the main theorem as follows. Firstly, since I = Ic holds (Condition (i)), observe that a predictor in (3) runs composition maps w & with w E W and $ E I Moreover, since the above second lemma (B) ensures that E Ieo can be represented by the composition of the projection onto X1, we can see that a predictor in runs o for some function class w e W", and hence, the bi-leveled optimization problem is expressed
min
Σ R eEEer
(10)
It is well- well-known that, assuming that runs all measurable functions,
à E arg min Rs Rs x1) - (((1)) E[y = æ1 |x 21) Pxę almost everywhere.
① E arg min (w xx)) + ((x)) = Prepx Prijx== Pxp almost everywhere.
1

---

Published in Transactions on Machine Learning Research (01/2024)
hold for any & under the least-square and cross-entropy losses respectively (Christmann & Steinwart, 2008 Example 2.6). Hence, ignoring the capability of W and the discussion of almost everywhere, we can see that
EX = xil se arg min
Σ eEE
(w 2x
(11)
or
PyTxX/= arg min
E eEEr.
Rw
(12)
12 and the first lemma (A), it concludes the main theorems intuitively, In the following section, we give the rigorous justification of the above rough proof sketch.
hold. Combining eq.s (11
4.2 Proof of Theorem
To prove the main theorem, we prepare two lemmas. Lemma 3. Let w X1 + y be the conditional expectation obtained by p' (y/z1); namely,
w' w11x1) EX = =x1]== == / y. p'''''(la (1lzz)dd.
Then,
w' o th E mig
min R (f).
1:X-y
Lemma 4. Any E I is represented as
for some continuous map V Xi + H.
Proof of Lemma 3 It suffices to prove the following statement:
For any F and (8"YY) E ((X', heEe, there exists (x vb) E ((X*, )EEE JeeE such that
xx f o yl1 dd yd (2, & -ll v1l -P (xx
(13)
Take arbitrary f E F and (xo, Va) E {(x*, Yo) Define (Xb, rb) E {(X*, Vo)eeee such that its distribution the direct product Pxp,y Px,: where Pxp,aa is the marginal distribution of Px on X, y and Px, is an arbitrary distribution on X2.
Then, the right-hand side of the inequality
13 is given by
/ [l(x)) vIIdPx
/ -1(()) (x,)
= / Pxo( (1((2),02) y/l PP (11,8).
Clearly, for any 2 E X2, the inequality
/ [1f((r) "/lpdxxx. (01,0)
|||Y X m1]
vll'dP (x1.
4The proof is essentially the same the one for Theorem in Rojas-Carull Carulla al. 2018 and Theorem in Toyota Fukumizu 2022.

---

Before providing a complete proof, we show a proof sketch of Lemma 4 to Published in Transactions on Machine Learning Research (01/2024)
holds, because the minimum of risk the least square loss is attained at the conditional expectation B[ya| Xf]. Hence, we obtain
/ - (lld) / Px, / [1(211 12) dP (xp
(x,( / |xiaaaaaa =1} /11 (x1
|B| = xix - yll yl (x1,1)
/ Pxy xg. (122) / |B| x x11 ulll x
= / || = xx)())) -
/ Il o x -
which concludes the proof. Here, the last equality is derived from the fact that the conditional expectation Xf = xi (x)] does not depend on & and corresponds to w! R.
Proof sketch of Lemma 4
make the flow of our proof casier to understand. First, we prove that E I can be represented as
+ = V 2h
(14)
by some map V X, + H. which is not restricted to a continuous map. Take arbitrary E I Then, since o E T I (Condition (i)), for any (Xa,y"), y"), (X5, Yb) E (o)e,,
Pyoeo(-) Prop(x))
and therefore, we have
(X110())) = (x))
(15)
14] by contradiction. Assune that there exist no maps that satisfy 14]. Then, there exist of E X1, I 2 E X such that
for any set N c y and Vr E X. We prove the statement
Φ x D
By utilizing of E Mi, 22 x2 E X2, we can construct (Xa, ya), (X*, yo) ((X"YY"))EE and c y which satisfy
(N
This contradicts the assumption 15 and we can conclude I can be represented as 14p The
continuity of v is easily derived from the continuity of 4 and we can conclude the proof.
Proof of Lemma First, we prove that & E I can be represented as
+ = v 2
(16)
by some map V X, H. which is not restricted to continuous map. We prove this statement by contradiction. Take t E I and assume that there exist no maps V which satisfy (16). Then, there exist o E Mi, 25 E X2 such that
+ (1) + (zf.
(17)
of (bxx] mt) *x az*) for any of X1, ag E X2, & depend only on the first component X1; hence, we can see that t TC can represented
Ni by some map v, which contradicts the assumption.

---

Published in Transactions on Machine Learning Research (01/2024)
y
supp (P ya)
N
supp (P xxxb,)b)
2
X
X²
xi
(Bon)
(90r)
X1
Figure 1: Supports of probability distributions Px yo and Px yu. The figure implies that Px (Ny,
Φ (())) + 0 and Pxa,ye (Ny- x ( = 0 (:: (x), & Φ (Φ*) 17 ) and that Pxo,y re -= (8- (T²)) #
and ( (4 (4x)) # O. These Eqs. lead us (NN ( * Py9-0(xo)((Ng-100) (o""
Fix y" E y with p' (v"1zt) > 0 and take an open neighborhood Ng- C y centered at y* which satisfies
0 0< / N.-
(dl 1.
Here, the existence of N is derived from the continuity of p' (+azf) (Condition (iv)).
Define two maps g* y X2 (i 1,2) by
(-(2) {
z (y Ny-)
( {
22
E E Ng-) clse )
22
(else)
Take two distributions "Y,,Y"), (x*, yd) E (x)e such that their distributions Px and Pxa.y coincide with
Px = Pxgiro @ Prixx & Pxt, Pxays = Pxyy=e & Pynxx @ Pxi-
Here,
Px, is a distribution on X1 where its p.d.f. coincides with a delta function dz; (x1) on zi,
the conditional p.d.f.s of Pxgiv=(-l) and Pxpivo(-1)) coincide with dy((w) (z2) and ()((()))2) respec- lively.
The supports of Pxaya and yr are visualized in Fig. 1 As Φ E Ico I (Condition (i)) and (X", "or,Y",, "xx,," E f(x", To
(N (0+) = (N, (T)),
(18)
where t² = (x... xz)) Let us compute Py (oy and Pyn (0*) to derive
Proporo) (h))= F Py (T)), which contradicts to the equality f18 We evaluate
(N (+)) (((x)))-=((0") x Ny.) Px-,y=($-1(6) x Ny-) (xxx))((*)))
Px ( (0-1(
Fig. illustrates the intuitive reason why PN 1')) (NV- derived, and hence, will help understand the following rigorous proof.
10

---

Published in Transactions on Machine Learning Research (01/2024)
by computing its numerator and denominator separately. First, the numerator is evaluated
(x- x N | J (0(()(022) p² (ylar) dz; (21) )drdy
S No- 20
dy So-1co (x2) p' (u|z)) on; (2) (d)d..
Noting that g²(y) = z. for Vy E Ny. and da; (21) x duz (22) = (41, 22), we obtain
Jkp dy S Jan 6g (12) p' (y/xz) d2g (x) = I No.-
dy Io- Ja-1(0n) (22(22) p' (ulxz) ozg (c) (r))do
= SSe No-
dy Sa-raa (11 (21, x2) pf (y)z1)da
Since (5) E x (0"), we have
Ssp dy So-coo, 10-1000000000555 (1111))011 p² (yllx)) SS, NNy of (ulli)ddy
Ny.
which leads us to the equality
( x Ny.) = JM,- p'(lliiddy.
Next, let us evaluate the denominator (x-1011(01)
Px- (2 ((-))))))) = P x V) = J (2(()((2)) p (v((=1) (x ddady
J dy I (o-100)) (22()) 1² (y) dm; )dx
NN dy I 1 (0'( p² (2)01) det (21)dr
+ *SS-N du So 1(00) 50( p' (y(m)) 82g (xi)de
J dy l Ja-10r) 80g (x22 p' (21z1) deg (21)
N,
*S--NN dy S (00)) (x.)(02) of (wlx1) deg (xi)dr
Jx dy S 1(00)) 8 (0)1 x22) (y/x1)dz No.
+/y -N-,-- dy 1o f0-10r)) 8(0t,33) p' (u)z) )dr.
Here, the fourth equality is derived from the facts that g2()) x5 for Vy Ng and g² x22 for Vy E Ny.. Noting that (x), x x-'( and (0),23") & Φ "f"(("),, we obtain
Px- ($ (o-'((()) = SS dy Jo-rar (109 (11, 12) pd (y(x))dr
Nu-
+ + dy Jo-1oo ((يtz)) (2i, x2) o' (ula )dx
= =In NN dy o' (ulo)) ( JN-N .
= =sn p'(llxiddy.
N.-
11

---

Published in Transactions on Machine Learning Research (01/2024)
Hence, we obtain
(x x Ny (N)-1)))))N))100)
Px- (x-))))())))
JN p' (ylxi)dy
JN p' (ylzi)dy
Next, let evaluate
Pocx)),ro(000 x Ny.) (a(x)(((*)) Px-(t-1(00)
Px),ys((- x Ny.)
(N
The numerator is evaluated as
1 (o-10) (g(()((2)) p² ())) dz; (on ddrdy (x) N,
/ dy So So-no ()2))))2)) p' (ulx1) do; (21)
Ng.
Noting that g²(()= = x²" for Vy E Ny, we obtain
Sng- dy fo-1oo (12())(22) p¹ (y/x1) 624(x1)do JxS- dy / Ja-r(ar) 60g((12) p² (vlæ1) deg (21 dd
[ dy [ Jo- (1) p¹ (y/ox)do
N₂-
[ dy 0 =0 O.
N₂-
Next, the denominator Pxa Px= is Here, the third equality is derived from (xf,25*) & Φ--((²"). evaluated as
- f (g2()( pf (z/x2) 82; (x1)dddy - dy L Jo-s(o)) (22)) x)) (x1(d
( (x-))) x v)
= Sm dy ( So-1(or) 604( (12) p (2)x1) oz; (x))
N₂.
JS-V. dy So dgc(n) o' (2)æ1) dny (x1)dx )dz
5S dy I J 8zg. (12) of (y/x1) dzg (x1)
Not
* dy Jo Ong (m2(222 -p' p' (1121)) dag (2 )d
= =|xg NNr
dy l/ - ((0) of (vl
+ $ dy Lo 1(00) 8607 (2) x2) p' (ula))
12

---

Published in Transactions on Machine Learning Research (01/2024)
Noting that x E + ( and (0f, 22*) & Φ ( we obtain
(x (Φ))) ((00)) / dy [ (-11 (111=1)(111 x2) p' (yla1)dz
NN.-
-f Nr- dy Sa-ror (1), x2) pt (g)xr )d
= =s Ny.
dy. 0+ f N-N,
dy dy (vl1z))
I p'(laitdy # 0.
Hence, we oblain
( x Ny (x-(1)1(0)) 0
PN
0.
JY-N,- p² (vlot)
Combing these results, we obtain
(N #0 0= PN
which contradicts the assumption
Projo(xe) = Py)
Because the continuity of V is trivial, conclude the proof.
Finally, prove Theorem
Proof of Theorem 1
Take
f E arg min seze. wewc
Σ (w
(19)
e9Er
Then, by Lemma 4 can represent f* as
=' w² z
for some continuous map w" : Xg + 7 Let us prove that w² o th E arg mig Roodd (f) by contradiction; assuming that w² o t. arg min, (J), we will derive w" tht & arg pere wewe Eecc,, R' o D), which contradicts to 19). We prove it by the following three steps.
Step 1 First, we prove that there exist a training domain e** E E and an open set N1 C X1 which satisfy
w (ax) a w (z1) for Va, E N
(20)
with Px.e (N1) > As w' Ri minimizes the o.o.d. risk (Lemma 3 we have
Rood (w" o t
min Rood (f) = Ro.o.d (w 2z1). f
By Lemma f can be represented by f w² 4 V for X1 + and + Y. Replacing V by wt" obtain the desirable statement.
13

---

Published in Transactions on Machine Learning Research (01/2024)
Here, the first inequality is derived from the assumption of proof by contradiction. Noting that Road. is maximum of risk among {(X*, Y*),, there exists (X*", y*) {(Xe, Ye)) such that
(= lly w² (x1))1 xi" .x-))
Rsd (w- o .
Rd Re (w' @ (- (=/t)>=/lbb wl1)-
(or))) yxr))
(21)
holds Since 21 is rewritten
/ flu A1ll-wi w² 22(x))1² - ll ww(n (x1))l dP X² yo- > 0,
dxxx'
we can see that
11 - w- (xt)112 Ily* - 2x]))11 (x'))= 00
for some (xi,8") E 24 x y. Since w and w! are continuous, taking sufficiently small 0, we have
lly" w (xx )l1 - Ill" w/ (a1))l > 0 for Vx1 E ot
(22)
where Ni; the e-ball centered at oj. Here, the continuity of wf derived from Condition (iv) in Theorem 1. Morcover, 22 leads us to the statement
w² (x1x) + w (xI) for Van N
By the condition (ii), No; ssspp(xxp" (x * for some ex E Eg. Take
"1" E No; U (uupppppxs Px.-.) def of supp
No M (ar &i Na, open neighborhood around x1 ( > 0} H.
Replacing af", if necessary, may assume that
i No; fon E E Na. open neighborhood around xı = (Pxp-))(Na) > =0.
(23)
Take open set Nı C NZ; which includes zt" Then, have
(x11 w (x)) for Vo N₁.
(24)
Observing that
21 E {a1 E X; Na, open neighborhood with 21 E No, = (Pxe))(NN, > 0}.
we have pxp (N1) > 0. It concludes the proof of Step 1.
Step 2 Next, we prove the inequality
Σ cEE.
Rw <
Σ
Re o (r
(25)
BEE
To derive the inequality, note that
@ min R (w X) + ((xx) = (x1) Pxp a.e.,
Note that is necessarily included training domains The inequality are proved Step (eq. 22
21 for training domain Etr
14

---

Published in Transactions on Machine Learning Research (01/2024)
equivalently,
① E arg min ( Th)) + (x E E& (0(xr) # w (x1) =
(26)
holds for any E & (Christmann & Steinwart
2008
Example 2.6). Taking the contraposition of the
implication from the left to right propositions in 26 we have
à satisfies Pxp (121 E Xi (0(11) + w' (I1) }) > = s & arg min R R" o xx)).
(27)
From 20 we have the inoquality
Pxp
(fr1 E Xi ( + w/ (x1) (M)
(28)
for some e* Etr and an open set N1 C N2 (27 and 28 lead to statement w' & arg min, R (woa and hence, we have the inequality
Rd (w o x))) min Rd (w (wo = R= (w" )
(29)
Moreover, since the conditional expectation w' minimizes the risk, we have
(w="" x)) Rno x²1)
(30)
for any E 3
29 and
30
lead us to the inequality
Σ
Rs (w xx) R (w" o x +
Σ
R2 o 2i)
eEErr (ct
R (w- o x ) +
Σ
Ru o x
R (w xx)++
+
Σ
Rs Re o xx)i==
E Rw o &)).).
eEEer (E'))
eEEur
Step 3 Finally, we prove w" o & & arg min sez WEW Eece.. (w o b), which contradicts to 19 By
the incquality 25 proved in Step 2, it suffices to prove that there cxist o' E I and wt E WC such that
w' o = wlo o Define Φ! = vr Φ. where the embedding vt Xi (= Rd) + H (= Rdx) ) is defined by
x
2
d
3
zd
E Rd
Here, we can define the embedding T since di & dx (Condition (iii). Noting that Py o1()) =
Projoxi (xe) Pytxx for any &, we can see that Φt ICo. Defining
x
2
x
dd 3
w²
E
j y.
d++1
can see that w &. = wl o Φt, Observing wl E W/C by Condition (iv), we can concludes w² & $ arg pize weweo Eeee, Re (w o b), which contradicts to 19
15

---

Published in Transactions on Machine Learning Research (01/2024)
4.3 Proof of Theorem 2
We prepare two lemmas..
Lemma 5. Let p² Xi + Py be the conditional p.d.f. of Pytpx; namely,
(v'x)), = = (ilx)
Then,
p' o th aag min R R.a.d. (f). f:X-PPy
Lemma 6. Any t Ih is represented as
for some continuous map V X1 + H.
Proof of Lemma The proof is essentially the the ones for Lemma hence, omit the
proof.
Proof of Lemma 6 First, we prove that & E I can be represented as
v xh
(31)
by some map V 2 - X, which is not restricted to continuous map. We prove this statement by contradiction in the same manner as the proof in Lemma Take & I 22,22" x2 E 22 such that
Then, there exist x E Xi
((11.0)) # (z),
Fix y" E y with p² (U"(xt) > O. Define two maps g¹ y N2 (i 1,2) by
g'()) = {
x2 (リ=
g²(v)) = {{
x2 2h2
(y (y= else )
xiz
else
Take two distributions (X",Y") (xx, E ('))e such that their distributions Pxa,y and yxa coincide with
Pxa.y= = Pxgiro & Pripx @ Px,, pxay== = Pxyire @ PY & Pxi
Here
Px, is a distribution on Xi where its p.d.f. coincides with a delta function dxz (a1) on xj,
the conditional p.d.fss of Pxçir= (-1y) and Pxgry=(Il) coincide with dy (x2) and (02((((22)) respec- tively.
Yo), (x", Y8) E ((X,, Yo)JEE, Since & Io = Ihe (Condition (i)) and ("),
Py (1 xT) ((v"") xT),
(32)
where x- (zf 24). Let us compute Py. ((1* (+) and (y)p)co))(1""(10) respectively. Same the proof in Lemma have the two equalities
Po(x-),y-((#"> (v*)) (x" and Po(x-)((99)) - (v*
which lead us to the equality
Py"pax")(("")0"" = Paxx ((t") (v"))
Po(x=)(f0b1)
=1
p'
16

---

Published in Transactions on Machine Learning Research (01/2024)
Similarly, we have
Py)p)())((("))1) and P
E
('(1=1) #0,
vEY-(0"
which lead us to the equality
(1)0)x))((1"")00) P (10*) {y*})
Po(xo)(((-))
Σ E (w)) p² (ylæ²)
Here, Ever-(o') p
E (ylx1) (ylz)) 0 derived by Condition (v). Combing these results, we obtain
Pr = + = (1000)))
which contradicts the assumption
Py- = Pr)
Because the continuity of v is trivial, we can conclude the proof.
Proof of Theorem 2 This is essentially the the one for Theorem 1, and hence, omit the proof.
5 Conclusions
In this paper, have proved that solution for the bi-leveled optimization problem (3) also minimizes o.o.d. risk (2) under four conditions in regression and classification cases, assuming that distributions on domains are the ones proposed in [Rojas-Crulla et al. (0018) and that models run all contimuous functions. Particularly, we have provided a sufficient condition on the training domains &,r and the dimension of the
feature space H. for the optimization problem (3 to minimize the o.o.d. risk.
Several challenges still exist. The first problem is the theoretical analysis of the optimization method for 33. To solve the challenging optimization problem 33 various optimization techniques have been proposed
(Arjovsky et al. 2019 Lin al. 2022a Zhou et al. 2022 and there has been little discussion about their
effectivenoss. For example, while Arjovsky et al. 2019 optimized (3] by minimizing
Σ R( + eEEτ.
V m 1.
R (w 1)
their effectiveness was evaluated only under specific SEMs (Rosenfeld et al. 2021 Thus, is important to
investigate this analysis under a more general case.
Second, should evaluate the o.o.d. performance of the bi-leveled optimization problem under the case
where the conditions in Theorems 1 and 2 are violated. Partieularly, noted in Section 2 condition Itr = I
does not generally hold. In such cases, for (4*, ",")"") arg min sers wewco Eeee, R (w 4),
Rood (w* Φ*) min Ro)
is not necessarily 0. The quantitative evaluation of the difference is crucial for future work.
Thirdly, we should investigate the feasibility of the condition I, = I, which is known to be an important and unsolved problem shared by all invariance- based methods (Arjovsky et al. 2019 Peters et al. 2016 Toyota Fukumizu 2022). As Condition (i) our main results, all methods based on invariances implicitly or explicitly assume that invariances among training domains correspond ones among all domains. As
discussed in Section 2 some sufficient, conditions under simple linear SEM setting have been found (Peters)
17

---

Published in Transactions on Machine Learning Research (01/2024)
et al. 2016 Arjovsky et al. 2019 but general theoretical results have not yet been established. This is also among our unsolved problems, and should be provided in further work.
Finally, extending our results to general domain sets beyond the case by Rojas- Carulla et al. (0018) is an important topic for future work. Invariant Risk Minimization (IRM) estimates the feature map that has the same conditional distribution Py=(((x*) among all domains e E E; other words, TRM framework assumes that domain set & has a feature map & such that (y-))(X.) are equal among all domains. Among domain sets that satisfy the property, the domain set by [Rojas-Cauula. projection Φ induces the same conditional independence Py-j--h In some cases, map that induces the same conditional distribution more complex function than the projection Φ so the relation between 3 and the o.o.d. risk such general domains beyond the case by PRoas-Caaulla et al. (2018) is should be investigated.
Carulla et al. (2018) is the simplest one; the
Acknowledgements
We thank Dr. Yano in the Institute of Statistical Mathematies for valuable discussions. The research was sup- ported by Grant-in-Ar for JSPS Fellows 20J213996, Grant-in-Aid for Research Activity Start-up 29K19966, JST CREST JPMJCR2015, and JSPS Grant-in-Aid for Transformative Research Areas (A) 22H05106.
References
Kartik Ahuja, Karthikeyan Shanmugam, Kush Varshney, and Amit Dhurandhar. Invariant Risk Minimiza- tion Games. In Proceedings of the 37th International Conference on Machine Learning, pp. 145-155, 2020.
Kartik Ahuja, Ethan Caballero, Dinghuai Zhang, Jean- Chan-Chhisttope Gagnon-Audet, Yoshua Bengio, Ioannis Mitliagkas, and Irina Rish. Invariance Principle Meets Information Bottleneck for Out-of-DDssrrbution Generalization. In Advances in Neural Information Processing Systems, 2021a.
Kartik Ahuja, Karthikeyan Shanmugam, and Amit Dhurandhar. Linear Regression Games: Convergence Guarantees to Approximate Out-of-Distribttio Solutions. In Proceedings of The 24th International Con- ference on Artificial Intelligence and Statistics, pp. 1270 1278, 2021b.
Ehab A. AlBadawy, Ashirbani Saha, and Maciej A. Mazurowski. Deep learning for segmentation of brain tumors: Impact of cross-institutional training and testing. Medical Physics, 45(3):1150 1158, 2018.
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant Risk Minimization.
arXiv:1907.08 028893, 2019.
Andrew R. Barron. Universal approximation bounds for superpositions of a sigmoidal function. IEEE Transactions on Information Theory, 39(3):930 (((3):930-945, 1993.
Sara Beery, Grant Van Horn, and Pietro Perona. Recognition in Terra Incognita. In Computer Vision ECCV 2018, volume 11220, pp. 472-489. 2018.
Shiyu Chang, Yang Zhang, Mo Yu, and Tommi Jaakkola. Invariant Rationalization. In Proceedings of the 2020.
37th International Conference on Machine Learning, pp.
Yonggiang Chen, Kaiwen Zhou, Yatao Bian, Binghui Xie, Bingzhe Wu, Yonggang Zhang, Ma Kaili, Han Yang, Peilin Zhao, Bo Han, and James Cheng, Pareto Invariant Risk Minimization: Towards Mitigating the Optimization Dilemma in Oft-of- Distributtn Generalization. In The Eleventh International Conference on Learning Represenlations, 2023.
Andreas Christmann and Ingo Steinwart. Support Vector Machines. Information Science and Statistics. Springer, 2008
George Cybenko. Approximation by superpositions of sigmoidal function. Mothematics of Control, Signals and Systems, 2(4):303 314, 1989.
18

---

Generalization. In International Conference on Learning Published in Transactions on Machine Learning Research (01/2024)
Will Douglas Heaven. Google's medical AI was super accurate in lab. Real life was different Slniicoorrddiabbbeeeeee-------dd
story. https://www.teccnnloogggg//
2020.
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer foodforward networks are universal approximalors. Newral Networks, 2(5):359 366, 1989.
Dongsung Huh and Avinash Baidya. The Missing Invariance Principle found the Reciprocal Twin of Invariant Risk Minimization. In Advances in Neural Information. Processing Systems, 2022.
Pritish Kamath, Akilesh Tangella, Danica Sutherland, and Nathan Srebro. Does Invariant Risk Minimization Capture Invariance? In Proceedings of The 24th International Conference on Artificial Intelligncce and Stalistics, pp. 4069-4077, 2021.
Masanori Koyama and Shoichiro Yamaguchi. When is invariance useful an -itto---Ditribttion General- ization problem ? Stfk::2008.08888) 2021.
David Krueger, Ethan Caballero, Joern- Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. Out-of-Distriuttion Generalization via Risk Extrapolation (REx). In Proceedings of the 38th International Conference Machine Learning, pp.
2021.
Yong Lin, Hanze Dong, Hao Wang, and Tong Zhang. Bayesian Invariant Risk Minimization. In IREE/CVF Conference on Computer Vision and Pattern Recognition, pp. 16000 16009, 2022a.
Yong Lin, Shengyu Zhu, Lu Tan, and Peng Cui. Zin: When and how to learn invariance without environment partition? In Advances in Neurul Information Processing Systems, 2022b.
Jiashuo Liu, Zheyuan Hu, Peng Cui, Bo Li, and Zheyan Shen. Heterogeneous Risk Minimization. In Proceedings of the 38th International Conference Machine Learning, pp. 6804-6814, 2021a.
Jiashuo Liu, Zheyuan Hu, Peng Cui, Bo Li, and Zheyan Shen. Kernelized Heterogeneous Risk Minimization, 2021b.
Chaochao Lu, Yuhuai Wu, José Miguel Hernández-Lobalo, and Bernhard Schölkopf. Invariant Causal Rep- resentation Learning for Representations, 2022.
Hrushikesh Mhaskar. Neural Networks for Optimal Approximation of Smooth and Analytic Functions. Newral Compulation, 8(1):164 1771161-77771111 1996.
Giambattista Parascandolo, Alexander Neitz, Antonio Orvieto, Luigi Gresele, and Bernhard Schölkopf. Learning explanations that are hard to vary. In International Conference on Learning Representations. 2022.
Christian S. Perone, Pedro Ballester, Rodrigo C. Barros, and Julien Cohen-Adad. Unsupervised domain adaptation for medical imaging segmentation with self-ensembling. Neurolmage, 194:1 11, 2019.
Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen. Causal inference by using invariant prediction: Identification and confidence intervals. Journal of the Royal Statistical Society: Scries B (Statistical Methodology), 78(5):947 1012, 2016.
Roman Pogodin, Namrata Deka, Yazhe Li, Danica J. Sutherland, Victor Veitch, and Arthur Gretton. Ef- ficient Conditionally Invariant Representation Learning. In The Eleventh, International Conference on Learning Representations, 2023.
Alexandre Rame, Corentin Dancette, and Matthicu Cord. Fishr: Invariant Gradient Variances for Out-of- Distribution Generalization. In Proceedings of the 39th International Conference on Machine Learning. pp. 18347-18377, 2022.
19

---

Published in Transactions on Machine Learning Research (01/2024)
Mateo Carula, Bernhard Schölkopf, Richard Turner, and Jonas Peters. Invariant Models for Causal Transfer Learning. Journal of Machine Learning Research, 19(36):1-34, 2018.
Elan Rosenfeld, Pradeep Kumar Ravikumar, and Andrej Risteski. The Risks of Invariant Risk Minimization. In International Conference on Learning Representations, 2021.
Janelle Shane. Do neural nets dream of electric sheep? AT Weirdness CeeronsseCCommeeuuhrrCCCnnnnnnn of
https:/ https::/wwwwaiiie http:/iiieiirdeee..ooddeeeeeeeee
2018.
Sho Sonoda and Noboru Murata. Neural network with unbounded activation functions is universal approx- imator. Applied and Computational Harmonic Analysis, 43(2):233 268, 2017.
Xiaoyu Tan, LIN Yong, Shengyu Zhu, Chao Qu, Xihe Qiu, Xu Yinghui, Peng Cui, and Yuan Qi. Provably invariant learning without domain information. In Proceedings of the 40th International Conference Machine Learning, 2023.
Shoji Toyota and Kenji Fukumizu. Invariance Learning based on Label Hierarchy. In Advances in Neural Information Processing Systems, 2022.
Xiao Zhou, Yong Lin, Weizhong Zhang, and Tong Zhang. Sparse Invariant Risk Minimization. In Proceedings of the 39th International Conference on Machine Learning, pp. 27222 27244, 2022.
20


Reviews:
[{"review_id": "NkYVTWSHg4", "reviewer": "Reviewer_g2Z9", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The results hold for the more general class of continuous functions than in previous works, which tended to focus on specific distributional assumptions on the training data or linearity of the function class.\n\nWeaknesses:\n- The presentation of the proofs and results could be improved. Please see my comments and questions below for specific requests.\n- Some of the assumptions are quite strong, e.g., Conditions (i) and (iii).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Deep networks are known to exploit spurious correlations in training data to make predictions, making their generalization capabilities to unseen or shifted domains challenging. Recently, Arjovsky et al. (2019) formulated an out-of-distribution (OOD) risk as a loss to improve the generalization of such networks, which consists of minimizing the worst-case risk over multiple domains. As this concerns solving a bi-level optimization problem (which is challenging to solve in general), they proposed Invariant Risk Minimization (IRM), which minimizes the average risk over all domains. The present work aims to provide a theoretical justification for their approach, by proving that solving the IRM problem achieves a predictor that also minimizes the OOD risk. The authors illustrate their results for both the least squares loss and cross-entropy loss, over a general class of predictors (the composition of two continuous functions), which is in contrast to previous works that focused on simpler classes of functions.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Since Theorem 1 is the main result and its proof is quite involved, I think it would be helpful to provide a proof sketch of the high-level ideas in the main body of the paper. This could be done before Section 3 or in the beginning of Section 4.\n- Condition (iii) feels odd to me as this requires the feature space to be larger than the data dimension. In typical deep networks, e.g., in image classification, the feature space is of much smaller dimension than the size of the images. Aside from the need to assume this in Step 3 of the proof of Theorem 1, is there a deeper reason why this condition is necessary?\n- Could the authors also discuss how their results differ in technical content to the nonlinear version of the Rojas-Carulla et. al’s Theorem 1 (Theorem 4 in Section A.1 of Rojas-Carulla et. al (2018))? This result is also in a nonlinear setting.\n- In the proof of Lemma 4, the conclusion $\\Phi(x_1^*, x_2^*) \\neq \\Phi(x_1^*, x_2^**)$ (eq. (14)) is crucial to the proof. Why is it the case that assuming there does not exist a $\\Psi$ such that (13) holds implies (14)?\n- Around equation (16), the next line should be due to Lemma 3, not 4.\n- In Theorem 1, shouldn’t the domain of $f$ be $\\mathcal{X}$, instead of $\\mathcal{H}$?\n- In equation (9), the argument of $R^{o.o.d.}$ should be $f$, instead of $p_{\\theta}$, since $f = p_{\\theta} \\circ \\Phi$.\n- Step 3, the inequality for $d_1$ and $d_{\\mathcal{H}}$ is backwards.\n- Some of the notation can be confusing, e.g., $N_{x_1}$ refers to a neighborhood around $x_1$, but points in $N_{x_1}$ are also denoted by $x_1$ (e.g., equation (20)).", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "mYCEi5e4pf", "reviewer": "Reviewer_JQnh", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- This paper is theoretically solid. The proofs are quite rigorous.\n- Most of the conditions are reasonable and can be achieved during small-scale datasets.\n\nWeaknesses:\n- There is no experimental validation to prove the proposed justification. Whether the theoretical results are applicable to real-world scenarios is still doubtful.\n- The significance of the conditions is not novel. There have already been several works showing the convergence property of hypothesis risk. How can the proposed justifications make further contributions is still not clear to me. Could the authors make some justifications?\n- Whether the least square loss and cross-entropy loss are still applicable to existing works is not explained. For example, CORAL, SWAD, and SparseIRM leverage external regularization terms to conduct invariant learning. Is the proposed theoretical results are still correct in these methods?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper theoretically proves that under certain conditions, IRM-based methods can indeed achieve invariant learning for OOD generalization. Specifically, two surrogate loss functions are considered: namely least square loss and cross-entropy loss, where the latter one is based on the results of the former one. As a result, some conditions for data distribution and hypothesis functions are achieved to justify the OOD generalization result of IRM.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please see the weaknesses.", "broader_impact_concerns": "The authors have addressed the broader impact in the main paper.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "oUgBDNskYv", "reviewer": "Reviewer_hemf", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n+ The paper analyzes IRM under more general distributional assumptions compared to prior work, without relying on specific structural equation models or linearity/Gaussianity assumptions. This expands the applicability of the theoretical results. \n+ A novel theoretical technique using proof by contradiction is introduced to characterize the form of invariant feature maps as compositions with the projection onto the invariant subspace. Sufficient conditions are provided on the training domains and feature space dimension for the bi-level optimization to minimize OOD risk. For instance, invariances captured by training domains must equal invariances over all domains.\n+ The theoretical results are proved in both regression and classification settings, for two common loss functions, MSE and Cross-Entropy. Continuous invariant feature maps and predictors are considered, making the function classes more realistic. Prior works often assumed simplified function classes like linear maps or variable selection, but this paper considers continuous functions which is more realistic for deep neural nets with universal approximation.\n\n## Weaknesses\n\n+ The analysis relies on the domain generation process proposed by Rojas-Carulla et al. (2018), where domains differ only in the conditional distribution of the spuriously correlated variables X2 given X1, while the causal mechanism P(Y|X1) is fixed. Expanding the theoretical results beyond this specific process to more diverse domain distributions would strengthen the conclusions.\n+ There is no analysis of the effectiveness of different optimization methods for solving the IRM bi-level problem in practice (e.g., IRMv1 from the IRM paper). Understanding optimization is key to successfully applying IRM in practice. \n+ Some sufficient conditions may be restrictive for real-world problems. For instance, the invariance equality between training and all domains may not hold generally. More analysis is needed for cases where the conditions are violated, and the feasibility of satisfying the conditions should be further discussed.\n+ Some numerical experiments would strengthen the paper, though not required. For instance, Rosenfeld et al. 2021 did some synthetic data experiments of IRM in the setup of their theoretical data model, and the results corroborated their theory well.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proves that under certain conditions, solutions to the bi-level IRM optimization problem minimize out-of-distribution risk, helping to justify IRM theoretically. Compared to prior work, the paper's analysis relies on more general distributional assumptions without specific structural equation models, and considers continuous invariant feature maps and predictors, making the theoretical results more widely applicable. Key technical contributions include using proof by contradiction to characterize invariant maps and providing sufficient conditions on training domains and feature space for minimizing OOD risk. The results are proved for both regression and classification settings. Overall, this significantly advances the theoretical understanding of IRM for OOD generalization under less restrictive assumptions than prior efforts.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See weaknesses.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper examines out-of-distribution (OOD) generalization performance of models trained via Invariant Risk Minimization (IRM), and proves that IRM obtains a predictor that minimizes the OOD risk under a broader class of distributions and functions than assumed by prior work.

The reviewers generally agreed that the results were rigorous and well-supported, and that the presentation was clear. While the overall setting is not novel and the specific conclusions are relatively narrow, the topic is salient and the paper will be of interested to a subset of the TMLR community.
Decision: Accept (poster)