Domain Generalization via Heckman-type Selection Models 
Authors: Hyungu Kahng
Hyungrok Do
Judy Zhong
Venue: ICLR 2023 notable top 25% (2023)
Keywords: Domain Generalization
Sample Selection
Bias Correction
Heckman

Abstract: The domain generalization (DG) setup considers the problem where models are trained on data sampled from multiple domains and evaluated on test domains unseen during training. In this paper, we formulate DG as a sample selection problem where each domain is sampled from a common underlying population through non-random sampling probabilities that correlate with both the features and the outcome. Under this setting, the fundamental iid assumption of the empirical risk minimization (ERM) is violated, so it often performs worse on test domains whose non-random sampling probabilities differ from the domains in the training dataset. We propose a Selection-Guided DG (SGDG) framework to learn the selection probability of each domain and the joint distribution of the outcome and domain selection variables. The proposed SGDG is domain generalizable as it intends to minimize the risk under the population distribution. We theoretically proved that, under certain regular conditions, SGDG can achieve smaller risk than ERM. Furthermore, we present a class of parametric SGDG (HeckmanDG) estimators applicable to continuous, binary, and multinomial outcomes. We also demonstrated its efficacy empirically through simulations and experiments on a set of benchmark datasets comparing with other well-known DG methods.

Published as a conference paper at ICLR 2023
DOMAIN GENERALIZATION VIA HECKMAN-TYPE SELECTION MODELS
Hyungu Kahng Hyungrok Do² Judy Zhong²
"ooraa University NYU School of Medicine
hgkahng@korea. ac. kr, (hyungrok. do, judy zhong org
ABSTRACT
The domain generalization (DG) setup considers the problem where models are trained on data sampled from multiple domains and evaluated on test domains unseen during training. In this paper, we formulate DG as a sample selection problem where each domain is sampled from a common underlying population through non-random sampling probabilities that correlate with both the features and the outcome. Under this setting, the fundamental iid assumption of the em- pirical risk minimization (ERM) is violated, so it often performs worse test domains whose non-random sampling probabilities differ from the domains in the training dataset. We propose Selection-Guidee DG (SGDG) framework learn the selection probability of each domain and the joint distribution of the outcome and domain selection variables. The proposed SGDG domain generalizable as it intends to minimize the risk under the population distribution. We theoretically prove that, under certain regular conditions, SGDG can achieve smaller risk than ERM. Furthermore, we present a class of parametric SGDG (HeckmanDG) es- timators applicable to continuous, binary, and multinomial outcomes. We also demonstrate its efficacy empirically through simulations and experiments set of benchmark datasets comparing with other well-known DG methods.
INTRODUCTION
In statistical learning theory, the standard assumption behind many supervised learning algorithms is that both training and test instances are independently and identically distributed (iid) according to the underlying data distribution (Vapnik, 1991). In other words, most statistical models assume that the training and test data are both random samples chosen randomly from the same pop- ulation. Unfortunately, this assumption is often violated in real-world applications rendering model performance to deteriorate on out-of-distribin Domain Generalization (DG) problem (Blanchard et al., 2011) 1) has gained particular attention, where it is assumed that learning systems have access to training data sampled from multiple domains, and the ultimate goal is to extrapolate new instances sampled from previously unseen test domains.
(OOD) test data (Koh et al., 2021). Recently, the
In this paper, we consider DG as a non-random sample selection problem. Let PxY represent the population data distribution, and S denote a binary random variable indicating whether subject is selected from the population into domain k. In random sampling process, P(S 1) is inde- pendent and identically distributed (iid). Under a non-random sample selection, the distribution of (X, Y) in domain A is conditional distribution of (X,Y) given S* = 1, which often does not equal to Pxy- Consequently, this leads distributional shifts across domains: Pky * P for k # j- Mathematically, distribution shifts across domains P may contain shifts in distributions of X (Pk, x covariate shift (Bickel et al., 2009)), and in the distributions of Y conditional concept shift (Moreno-Torres al., 2012)). We present a graphical model in Figure conceptu- ally illustrate the sources of distribution shifts, assuming the existence of latent factors confounding the relationship between X, Y. and domain (S*). In Figure 1, C1 represents unobserved latent factors that correlate with X and S* resulting in covariate shift. C2 correlates with X, Y, and S si- multaneously, entailing both covariate and concept shifts. The goal of DG is estimate the domain generalizable (agnostic) edge ʃ X + y in the presence of the two types of latent confounders.
X (P), (
"This work mostly done while the author was with NYU.

---

Published as a conference paper at ICLR 2023
Figure 1: A graphical model illustrating the source of dis- tributional shifts. X: covariates, Y: outcome, s" domain. Ci represents latent factors that correlate with and s*. resulting in covariate shift. C2 correlates with X, Y, and s*, entailing both covariate and concept shifts. Our goal is to estimate the do- main generalizable (agnostic) edge f X + y in the presence of the two types of latent confounders.
x
C₁
f
Y
C2
St
The vast majority of DG methods are developed to identify f that is robust to C1. However in practice, C2 type of confounders often exist which make P(S 1) related to both X and Y. For example, when we train a model to predict tumor status (Y) from histological images (X) using patients from different hospitals (S*), there may be variations in X due to inconsistent acquisition processes such as staining differences (C1) across hospitals, and differences in patient characteristics such age, gender, race, and discase severity (C2) that correlate with hospital, covariates and the outcome. As a result, a model trained in an oncology specialist hospital may not be generalizable to hospital serving veterans. Similarly, when train model to predict wealth index (Y) from satellite images (X) taken from different countries (S*), ) there may be latent factors such as the economic status (C2) correlating with X, Y and domains simultaneously. Therefore, model trained on one country may not perform well in another country with a different rural/urban proportion or economic status (Koh al., 2021).
In this paper, we propose a new class of Selection Guided Domain Generalization (SGDG) models to first estimate the selection probability that an instance is sampled into a training domain, and then use the joint distribution of the outcome Y and selection S' to learn a domain generalizable model. In particular, SGDG built on Heckman's 's bias correction framework (Heckman, 1979) which is very powerful tool to learn an unbiased model from non-randomly selected samples in the presence of both C₁ and C2 confounders. The unique contributions of this paper are summarized as follows:
To the best of our knowledge, we are the first paper to formulate the DG problem using a non-random sample selection framework, and to propose a Selection Guided Domain Generalization (SGDG) method under this framework.
We present class of parametric SGDG (HeckmanDG) estimators applicable to continuous, binary, and multinomial outcomes
We demonstrate the efficacy of our method both theoretically and empirically on simulated data and four challenging benchmarks.
RELATED WORK
Domain Generalization. DG has been studied under various contexts. Many studies are devoted to learning dlmain--nvariann features which are discriminative and independent of the domain, such as kernel-based methods (Muandet et al., 2013), matching moments (Sun & Saenko, 2016), adversarial learning (Ganin et al., 2016; Deng et al., 2020), entropy regularization (Zhao al., 2020), and contrastive learning (Motiian et al., 2017; Kim et al., 2021). ). Other works exploit invariant causal effects across domains (Arjovsky al., 2019; Ahuja et al., 2020; Rosenfeld et al., 2021). Another family of robust optimization methods seek to minimize the worst-case error (Sagawa et al., 2020; Xie et al., 2020; Krueger al., 2021). More recently, other prominent directions of methods improve DG by model averaging (Cha et al., 2021; Arpit et al., 2022), gradient matching (Shi et al., 2022), meta learning (Li et al., 2018), data augmentation (Robey et al., 2021), 1),, and generating novel domains (Zhou et al., 2020).
Sample Selection Bias Correction. Zadrozny (2004) formalized sample selection bias in machine learning terms and presented bias correction method when selection only depends on the input features. Cortes et al. (2008) proposed sample reweighting approach to tackle the same problem but assumed the availability of additional data drawn from the true population. Du & Wu (2021) proposed a framework for robust and fair learning under biased sample selection, but assumes con- ditional independence of Y and S given X.
code available: https: htt.ub.com/hgaaaggg//ee generall zat i on- 1nn-ingtniinn ng
2

---

Published as a conference paper at ICLR 2023
3 SELECTION MODEL-GUIDED DOMAIN GENERALIZATION
3.1 DOMAIN GENERALIZATION
Suppose there exists L distinct but relevant domains and let S = (S) (S)..... Ps denote binary random vector that indicates the domain membership where s = implies belonging to domain k. Let Pxy represent the population data distribution, and each domain's 's data distribution
be a conditional distribution of the population distribution given S = 1, i.e., Pxx Pxy|s*=
Assumption (Mutually Exclusive Domain Membership). We assume that if S 1, then S = for all j + k so that an instance can belong to one and only one domain. Assumption (Independent Domain Sampling Processes). We assume that S S for all j # k.
In the supervised domain generalization context, we are allowed to observe the joint distribution of X and Y, Pxx for K out of I. domains, and refer the K domains source or training domains observed during the training phase. The remaining L K domains are referred to as target or testing domains whom we may observe the testing phase. Under this setting, we aim to learn a prediction model that generally performs well on both source and target domains. We formalize the problem as follows.
Definition (Domain Generalization (Blanchard et al., 2011)). Domain generalization refers to as the problem of learning f X y that has the minimum expected loss across all possible domains, which can be further summarized as the following optimization problem:
(() =
min FEF (XY))xX
E
where E y E 0}) is loss function and F is hypothesis set. We note that some other papers have considered the same problem setting (Muandet et al., 2013; Deshmukh al., 2019; Blanchard et al., 2021).
We first introduce a proposition that claims the equivalence of the domain generalization and the risk minimization under the population distribution:
Proposition (Equivalence of Domain Generalization and Population Risk Minimization). Problem (1) is equivalent to the risk minimization under the population distribution Pxy. That is
= min (EF(XYY)-XxY [el(()))))]
min
E.
FEF (XYYP''
which straightforwardly follows from the law of total expectation. We define fERM the minimizer the the population risk minimization problem, as well as the best (hypothhetcaal) model for the domain generalization problem.
Proposition important as establishes that the best model for the domain generalization problem should minimize the risk under the population distribution Pxy. However, in practice we have no access to the population data distribution and are only given the source domains during model training. A naive learning model is minimizing the risk under the source domain data distribution:
Definition 2 (Source Domain Risk Minimization). Given the source domains whose data distribu- tions are xxx Pxy = for k = I,... K, we refer to the following learning problem as the source domain risk minimization problem
()) =
min JEF
E (XYY)P
(3)
We denote its minimizer as FSDRM- The empirical form of JSDRM is denoted as empirical risk mini- mization (ERM).
The generalization performance of fSDRM depends on how well the source domains represent the population (or the target domains). If the source domains well approximate the population distri- bution, then the generalization performance of fSURM will be sufficiently close to that of fPRM- On the other hand, if Pxxy # Pxy, one cannot guarantee that models trained on the source domains to effectively generalize to the population. Therefore, is necessary to model the selection probability to bridge the gap between x and Pxy.
3

---

Published as a conference paper at ICLR 2023
3.2 SELECTION MODEL-GUIDEL DOMAIN GENERALIZATION
Now we derive our selection model-guided domain generalization problem. starting from decompos- ing the objective function of the domain generalization problem by introducing domain selection model g = [gk: : X + [0, 1], k = 1,... L) which predicts the selection probabilities of an instance being observed in domain k.
Assumption 3 (Decomposable Loss Function). We assume that the loss function & can be decom- posed into Iwo components one exclusively about the selection model and the other involving both:
L
((f(X),
I(s 1)A(f(()) =1)
where lg : 10, 11 1) - R is the loss function for learning the domain selection models and A: : (y x [0, 114) x (V' x f0, 10- + IR+ is the joint loss function, which assigns to the prediction and domain selection models pair of true outcome Y and domain membership indicator S.
For example, the negative log-likelihood under the probabilistic framework satisfies Assumption 3: if we let ((f(X), Y) = - log p(Y|X), then A - log (Y|],, S = 1) and l, log p(S 1 I)), respectively. We provide the full derivation in Appendix A. and present a specific parametric form in Section 4.
Under Assumption 3, the population risk can be expanded as follows:
E (XXY)XPY
[(((())))) =
E
A(((X), = +1)+
p(s = 1)
K
Σ
E (XYipper
[AU((X) | =1) P
Σ E xxpp
((()()), pls = 1)
Σ
((X =1) P(S* = 1).
(4)
k=k+1 (X,))
Based on the expansion, we introduce our selection model gooee-ggided domain generalization problem. Definition 3 (Selection Model-Guided Domain Generalization (SGDG)). Given access to the data distribution of K source domains, P for A = 1,.... K and to the unlabeled data distribution from all L domains, Px|s==1 for k = =1..... 1....,.. L. We define the selection model-guided domain generalization problem as a joint learning problem of f : X + y and fgr
L
min SEF.9EG
E
(XX = p(s = [1]
1 k=
E Xapx
11s = =1),
(5)
where and G are hypothesis sets.
The SGDG problem is a minimization of the expected loss of the prediction model f under the joint distribution of X and Y of the source domains, and the expected loss of the domain selection model under the joint distribution of X and S. Under this formulation, the domain selection model guides f to be corrected through A, considering the probability of being drawn from certain do- mains. In Section 4, we will introduce specific forms of this problem under parametric assumptions. Theorem (Performance Improvement of SGDG over SDRM). Let fSDRM and fSGDG be defined as in Definition 2 and 3. fSGDG has lower risk than that of JSDRM. That is,
E
[C(]sCDG (X),caa(X)))) (X), (X,Y)~xx [(([sDRM(X)))))
E.
(6)
(x
which implies SGDG is expected to show better generalization performance compared to SDRM. By Proposition 1, SGDG is expected to perform better than SDRM for the domain generalization problem equivalently. The sketch of proof of this theorem is that SGDG performs as well SDRM on the source domains and offers performance improvement on the target domains. Proof of this theorem is given in Appendix A.2.
In reality, we may have no access to the distribution P for all domains (k: = 1,.... L) but only have access to the source domains' distributions (k = 1..... ,.). In such cases, we use P for k 1,... K to learn the selection model g as in Problem (7).
4

---

Published as a conference paper at ICLR 2023
Assumption 4. Let g" be the optimal selection model that has the minimum expected Es.
g = 00i .2L argmin
E Ex (
11) =1).
We assume that the first K coordinates (gi, (9) minimize Σtn E Exppx (x [[]
Conceptually, Assumption 4 means that the same selection models gk can be learned by contrasting
domain & to the remaining domains in the training data (K) {(k) as by contrasting domain & to the remaining domains in the population (I (k.)). Under Assumption 4, Problem (5) reduces
((XXX)) = - = =1+ km xmpp
5c =).
min SEF-9ES
(X.Y)~
(7)
HECKMAN-TYPE SELECTION-GUDDER DOMAIN GENERALIZATION
In this chapter, we present parametric models for A embodying fSGDG presented in the previous section. The essence of A is to model the conditional distribution of P(Y|X, S = 1). Consider the setting where we are given training data from multiple source domains in the form of D { 8, Yi J where s; [ss1, 812,.... ... 8ak: ... sik] {0, 1) K is binary vector indicating domain membership. By formulating the loss function as the negative log- liglikelihood, the empirical form of Equation 7 becomes
K
Si log p(u: 8ik = 1) + 9ik log p(s = +1x.)+((11110 sog 0|z:) (8)
min
Heckman (1979) proposed to model the joint distribution P(Y, SIX) of the selection latent variable S, where S = I[S 0]. and the continuous outcome Y via bivariate normal distribution with correlation coefficient and the mean as a linear functions of X. Building upon his work, we make the following assumption the joint distribution of the outcome and the selection variables.
Assumption (Joint Distribution of (or latent variables Y) and Š). Let be the outcome variable and 5t be the latent continuous variable where S = I[Š* > 0] for each domain k. We assume that Y and S (S) .........."" are jointly distributed as multivariate normal distribution with mean (f(X),9x())) and correlation coefficients (px { given X.
For binary and multinomial outcomes, this assumption on Y as the latent continuous variables underlying the observed outcomes. Under this assumption, = {9x | can be modeled as a set of independent probit models. three cases, the joint log-likelihoods have closed forms. Hence- forth, we refer to the specific parametric form of SSGDG as the Heckman-type DG (HeckmanDG) estimator.
Definition 4 (Heckman- Typp Domain Generalization Estimator). We formulate HeckmanDG as joint learning problem of and g = {9+ HK with the following learning objective:
K
sax - log (w.)) + (1- log (- (xx)) H (9)
1.
where ((-) is the cumulative distribution function of the standard normal distribution d, such that ( (x.) ) P( = 1|X = æx) is the selection probability w.r. domain k. and (1-g((<))) P(S 0|X = 2.). Meanwhile, A(f(x;;) 9t. (x:)) Us, Sski E) the conditional negative log prob- ability of ys given s i.e., log p(y: Sik 1.).
The specific form of A and miscellaneous model parameters Σ depends the prediction task (either continuous-valuee binary, or multinomial outcome prediction). For example, the objective function (9) for the continuous outcome is listed below in Equation 10, where assume Y =[(X)= + e,
5

---

Published as a conference paper at ICLR 2023
Figure 2: Schematic
Denaais ad acaccio modd
MAt
Step Learn
overview of HeckmanDG, with K training do- mains. Step 1) We learn g [Wg, (
lag ((u(ss))) so(x..)
Oeatue extineto))
to predict the selection probabilities for each training domain. Step 2) The estimated domain selection model 9 guides f = wy o ys to be cor- rected by considering the selection probability of instances being drawn from the training domains.
oty
f(mi)
Ouuaaa (dsadod)
Shop Loar o Pp.
Outoomec predictioos rookd.
S* [[g (X) + 7h > OJ, and m] ~ No [1, 0; Pk]). The details of derivation and for other types of outcomes can be found in Appendix A.3.
log & pM 31-((z)) ((x,) Sik
+ 9n (121) + log 61 ((mm)(((;)
min 1.9
+ (1 - 8.t)) log Ob (x (mx)).
a
(10)
Ultimately, f is the outcome prediction model of primary interest. For any input x from unseen domains, the HeckmanDG prediction is f(x).
The proposed HeckmanDG differentiates from Heckman's bias correction method in the following perspectives. First, it allows and models multiple domain-specific sample selection mechanisms. Second, by utilizing the multiple domains in the training data, HeckmanDG does need auxiliary data (features of instances from the target population but not sampled in the domain), which is required input for the original Heckman model. Third, HeckmanDG alleviates the linear assump- tions in Heckman's 's bias correction model to allow flexible forms for both prediction and selection functions including neural networks.
OPTIMIZATION
In this section, we consider the hypothesis sets FF and 9 are neural networks, and present an efficient algorithm to optimize Equation 9. Specifically, we propose two-step approach that primarily trains g to optimum (Step 1), then updates the remaining parameters of and Σ (Step 2). Step we learn the selection model parameters based on the following objective function:
g argmin ΣΣ
sik log ligh (xx)) + (1 s.k) log 2 (xx)))
(1I)
which essentially learns predict the source domain memberships for training instances. In Step 2, we freeze the selection model ÿ. and learn the remaining parameters of and Σ. The proposed optimization algorithm in part motivated by Heckman's two-step estimator for bias correction, which was devised as means of avoiding the non-linearity of estimating both selection and out- come equations simultancously (Heckman, 1979). We provide pseudocode describing the overall optimization procedure in Algorithm 1 in Appendix A.4, and an ablation study in Section 7 to sup- port the necessity of two-step optimization.
Neural Network Architecture. We use a common feature extractor Py X + 2, for the domain
selection models g {9u ) which only differ in the last linear predictors Wg.k Z, - [0, 1]. thus 9 = Wg o gg This prevents the number of parameters of g from growing with the number of training domains K, and reduces computational complexity by requiring only a single forward pass (.). Similarly, we define f o where ypy X' - Z, is the feature extractor of the outcome prediction model and w 2, + y is the final linear predictor. In general, 99 and f are allowed to have different neural architectures with different numbers of parameters. However, we found that simply using the same architecture (but learning different parameters) works well in practice. Therefore, in all our experiments, HcckmanDG has roughly twice as much trainable parameters as other comparative methods. An overview of our neural network architecture is provided in Figure 2.

---

Published as a conference paper at ICLR 2023
Train
ID
0OD
Random
ERM (oracle)
1.57 (0.46)
1.30 (0.49)
1.30 (0.40)
1.01 (0.04)
ERM
0.84 (0.09)
0.89 (0.12)
14.65 (6.00)
6.84 (1.83)
IRM
0.83 (0.09)
0.88 (0.11)
16.45 (6.40)
7.42 (1.86)
GroupDRO
0.91 (0.39)
0.88 (0.12)
15.83 (5.22)
7.33 (1.67)
VREx
0.85 (0.08)
0.89 (0.13)
14.47 (4.91)
6.96 (1.68)
HeckmanDG (ours)
1.40 (0.24)
1.37 (0.29)
8.71 (3.96) 3.70 (0.96)
Table 1: Predictive performance on simulated data, measured in terms of the mean squared error (lower is better) averaged over 30 trials. Standard deviations are given in parentheses.
EXPERIMENTS
Simulation. We simulate a linear regression problem to assess HeckmanDG's predictive perfor- mance. We simulate two covariates (X) and two training domains (K 2) based on the following setting of domain selection and outcome mechanisms:
s
Tla + oz X2 + Mh > O).
Y 1 155x, + 3X2 +E
X1 X2
~ ~N(O 12), and
mi
~ N
(
1000
Pko o²
(12)
Pko
where we assume normal prior over the selection coefficient ag N(Pag o2 in order to implic- itly control the similarity between domains by differing its parameters. We consider a true population of 100000 instances, from which we sample data for each domain with ni Uniform In each trial, we sample the selection coefficient a5 ~ (5, (²) for the two training domains and a held-out in-distribution (ID) test set (similar to the training domains), and af ~ (-(5,, ((-5,3²) 32) for another held-out out-o--distribution (OOD) test set (dissimilar the training domains). We also simulate a random test set from the population. We assumed
2000).
0.8 and a = 1.
Simulation Results. We observed that HeckmanDG not only outperforms ERM, but also other DG methods including IRM (Arjovsky et al., 2019), GroupDRO (Sagawa et al., 2020), and VREx (Krueger et al., 2021) (Table 1). Note that 'ERM (oracle)' is trained on iid training data, which serves as theoretical lower bound on model performance. We highlight that not only does HeckmanDGi perform well on the random test set, but also on the OOD test set. In contrast, other methods tend to fit well to the train domains (II) performance is high), but generalize poorly to random and OOD test domains.
Benchmark Datasets. To further demonstrate the effectiveness of HeckmanDG on high- dimensional data regimes, we conducted experiments on four datasets from the WILDS benchmark (Koh et al., 2021): 1) CAMELYON 2) POVERTYMAP, 3) IWILDCAM, and 4) RxRx1. We used the same neural network architecture for selection and outcome feature extractors i and y which followed by linear predictors wy and wy. For the domain selection model, tuned hyperparam- eters to obtain the best domain selection model returning the highest macro score on the training data, which in all cases achieve near-perfe peerect accuracy. For the outcome prediction model, we ad- hered to the official guidelines and used the OOD validation set provided in the WILDS repository for hyprraaraeeter tuning and model selection based on the recommended metries. Detailed de- scriptions of dataset statistics are presented in Table 5 of Appendix A.4. Details on hyperparameters and model training are presented in Table of Appendix A.4.
Benchmark Datasets Results. We summarized the results on the WILDS benchmark in Tables 2,3, and 4. All methods apply the same network architectures based on the WILDS guideline for fair comparison. Also, we point out that we exclude methods that deviate from the DG setting such as those centered on test time adaptation and using additional unlabeled data. We observed that HeckmanDG outperforms other methods two out of four datasets, while performing on par with other methods on the remaining two. We make the following key observations. First, HeckmanDG often robustly performs the test domain although it may not generate the best performance the validation domains. For example, on the CAMELYON dataset, the performance gap between the validation and test datasets for HeckmanDG is 3.3, which is substantially smaller than 14.1 1.1 of

---

Published as a conference paper at ICLR 2023
Method
Validation
Test
ERM (scratch)
84.9 (3.1)
70.8 (7.2)
ERM (ImageNet)
91.3 (0.2)
84.2 (1.7)
CORAL
86.2(1.4)
59.5 (7.7)
IRM
86.2 (1.4)
64.2 (8.1)
GroupDRO
85.5 (2.4)
68.4 (7.3)
VREx
(1.3)
71.5 (8.3)
LISA
81.8 (1.3)
77.1 (6.5)
Fish
82.5 (1.2)
79.5 (6.0)
SWAD
88.1 (1.5)
83.9 (0.9)
122-OT
86.3 (3.4)
77.5 77.5 (6.7)
HeckmanDG (ours)
90.6 (2.4)
87.3 (2.4)
Table 2: CAMELYON 7: We report predictive performance measured in terms of average accuracy on both the OOD validation and test set. Standard deviation across 10 replicates are given in parentheses. 'ERM (scratch)' trained from random imitial parameters, whereas we also report 'ERM (ImageNet)" trained from ImageNet-pretrainc weights (Russakovsky et al., 2015).
Average
Worst Group
Method
Validation
Test
Validation
Test
ERM
0.80 (0.04)
0.78 (0.03)
0.51 (0.06)
0.45 (0.06)
CORAL
0.80 (0.04)
0.77 (0.05)
0.52 (0.06)
0.44 (0.06)
IRM
0.81 (0.03)
0.77 (0.05)
0.53 (0.05)
0.43 (0.07)
GroupDRO
0.78 (0.05)
0.75 (0.07)
(0.04)
0.39 (0.06)
DANN
0.77 (0.04)
0.69 (0.04)
0.44 (0.11)
0.33 (0.10)
Fish
0.81 (0.01)
0.81 (0.01)
SWAD
0.78 (0.03)
0.77 (0.04)
0.48 (0.09)
0.45 (0.11)
HeckmanDG (ours)
0.81 (0.03) 0.81 (0.03)
0.53 (0.06)
0.51 (0.04)
Table 3: POVERTYMAP We report predictive performance measured in terms of both the average and worst- group Pearson correlation coefficient, on both the OOD validation test sets. Standard deviation across replicates are given in parentheses. We use the original 5 folds provided in the WILDS repository. We do not report worst-grow group performance for 'Fish', since it has not been reported in Shi et al. (2022).
'ERM (scratch)' and 7.1 of 'ERM (ImageNet)'. This pattern is similar what observed in the simulation studies. Because HeckmanDG is designed to predict for the underlying population, may not produce the best prediction performances for the (non-randomly selected) source domains. Although we do not know the specific selection probability of the test domain, HeckmanDG is often more robust for the testing domains by optimizing for the population distribution. Second, HeckmanDG effectively works for both the domain shift and the subpopulation shift problems. This can be observed by looking either the average or worst-group performance on the POVERTYMAP dataset (i.e., 0.51). This again supports the robustness of HeckmanDG against range of non- random selection probabilities the testing domains.
ANALYSIS
Necessity of two-step optimization. To demonstrate the necessity and effectiveness of the pro- posed two-step optimization, we performed an ablation study on the POVERTYMAP dataset. For comparison, we trained g and f in one-step to simultaneously minimize Equation 8. We observed that one-step obtained comparable performances on training domains (Figure 3a), but yielded worse test performance the OOD domains (Figure 3b). We believe that suboptimal ğ may mislead to deficient solution in the one-step optimization (more details in Appendix A.6).
Relationship between the performance of g and f. To further investigate how the performance of a domain selection model influences the final outcome prediction, we take intermediate snapshots for g. and learned / based each snapshot. On all five data split folds, always observed positive
8

---

Published as a conference paper at ICLR 2023
RxRx1
(b) IWILDCAM
Method
Validation
Test
Method
Validation
Test
ERM
19.4 (0.2)
29.9 (0.4)
ERM
37.4 (1.3)
31.3
CORAL
18.5 (0.4)
28.4 (0.3)
CORAL
37.0 (1.2)
32.8 (0.1)
IRM
15.2 (0.4)
(1.1)
IRM
20.2 (7.6)
15.1 (4.9)
GroupDRO
23.0 (0.3)
GroupDRO
26.3 (0.2)
23.9 (2.1)
LISA
20.1 (0.4)
31.9 (1.0)
DANN
31.9 (1.4)
Fish
(0.6)
10.1 (1.5)
Fish
25.8 (0.5)
24.2 (0.9)
SWAD
14.2 (0.5)
22.9 (0.7)
SWAD
31.6 (0.2)
29.1 (0.1)
L2A-OT
17.5 (0.3)
27.8 (0.9)
L2A-OT
22.8 (2.9)
18.1 (3.2)
HeckmanDG (ours)
20.5 (0.7)
32.1 (0.8)
HeckmanDG (ours)
34.5 (0.9)
31.8 (0.3)
Table 4: A summary of results on RxRxI and IWILDCAM. Note that the OOD validation macro FI
score for DANN was not reported since it has not been provided in Sagawa et al. (2022).
0.8
Drs step
Optinization
0.8
Sito (proosse))
0.85
0.6
Outcome (Pearson)
0.80
Pcarson
Pearson
0.4
0.6
0.70
one step lisil wett gras))
0.4
0.65
Step /Pain awemael
Average
Worst Group
Epochs
Eva luation
Selection (Macro
(a)
(b)
(c)
Figure 3: Analysis based on the POVERTYMAF dataset. 3a) Pearson correlation measured on the training data for HeckmanDDG learned by one-step (gray) and two-step (bluc, proposed) optimization. 3h) Comparison between seestep (gray) and two-step (blue, proposed) optimization, where performance is averaged across all five folds provided in the WILDS repository. 3c) Relationship between the performance of domain selection and outcome prediction (colors indicate different folds).
correlation between the performances of ğ and (Figure 3c), demonstrating that well-performing selection model necessary to correctly guide the outcome prediction model.
CONCLUSIONS
We propose a Selection Guided Domain Generalization (SGDG) framework, in which we formulate domain generalization prediction model f and the domain selection model g to achieve generalization on the true popula- tion. DG is challenging problem the particular structure of the distribution shift in the testing domains is unknown. In the presence of this uncertainty, we propose and theoretically justify the ob- jective of minimizing the risk targeting the population distribution through SGDG. Furthermore, have provided a set of Heckman-type SGDG estimators for various outcome types under parametric assumptions. Although it is still open question if single general- purposs training algorithm can produce models that do well on all of the DG datasets (Koh et al., 2021), we observed robust performances of the proposed HeckmanDG on four benchmark datasets.
non-random sample selection problem and propose to jointly learn the
Note that SGDG can naturally utilize all domains in the training data, including the (outcome) labeled and unlabeled domains. The unlabeled domains will contribute the estimation of ğ in Equation 9, which indirectly improves generalization performance as we showed in Section 7. An intriguing direction for future research to explore whether we can improve SGDG performances by adapting ((X) for domains in the test data. In that way, we may further refine the prediction for the unseen domains guided by their similarities to the source domains.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENT
We thank the ICLR reviewers for their suggestions on improving the original version of this pa- per. This work was supported in part by the National Institutes of Health under awards NIH R01- LM013344, R01 R01A4644467, RO1 AG065333.,, and R01-AG065330-0251 by NYU Center for the Study of Asian American Health under the NIH/NIMHD grant award #U54MD000538 (HD).
REFERENCES
Kartik Ahuja, Karthikeyan Shanmugam, Kush Varshney, and Amit Dhurandhar. Invariant risk mini- mization games. In Proceedings of International Conference on Machine Learning, pp. 145-155. PMLR, 2020.
Martin Arjovsky, Léon Bottou, Ishaan Guirajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint 1eX.i:1907.0222 2019.
Devansh Arpit, Huan Wang, Yingbo Zhou, and Caiming Xiong. Ensemble of averages: Improving model selection and boosting performance in domain generalization. In Proceedings of Advances in Neural Information Processing Systems. Curran Associates, Inc., 2022.
Steffen Bickel, Michael Brückner, and Tobias Scheffer. Discriminative learning under covariate shift. Journal of Machine Learning Research, 10(9), 2009.
Gilles Blanchard, Gyemin Lee, and Clayton Scott. Generalizing from several related classification tasks to a new unlabeled sample. In Proceedings of Advances in Neural Information Processing Systems. Curran Associates, Inc., 2011.
Gilles Blanchard, Aniket Anand Deshmukh, Urun Dogan, Gyemin Lee, and Clayton Scott. Domain generalization by marginal transfer learning. Journal Machine Learning Research, 22(2):1-55,
2021.
Junbum Cha, Sanghyuk Chun, Kyungjae Lee, Han-Cheol Cho, Seunghyun Park, Yunsung Lee, and Sungrae Park. Swad: Domain generalization by seeking flat minima. In Proceedings of Advances in Neural Information Processing Systems, pp. 22405-22418, Curran Associates, Inc., 2021.
Corinna Cortes, Mehryar Mohri, Michael Riley, and Afshin Rostamizadeh. Sample selection bias correction theory. In Proceedings of the 19th International Conference Algorithmic Learning Theory, pp. 38-53. Springer- Verlag, 2008
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.1 14478, 2020.
Aniket Anand Deshmukh, Yunwen Lei, Srinagesh Sharma, Urun Dogan, James W Cutler, and Clay- ton Scott. A generalization error bound for multi-class domain generalization. arXiv preprint 1.Xiv:1905.10392 2019.
Wei Du and Xintao Wu. Fair and robust classification under sample selection bias. In Proceedings of the 30th ACM International Conference on Information & Knowledge Management, pp. 2999- 3003, 2021.
Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François training of neural networks.
Laviolette, Mario March, and Victor Lempitsky. Journal of Machine Learning Research, 17(59):1-35,, 2016.
Vassilis A Hajivassiliou and Paul Ruud. Classical estimation methods for ldv models using sim- ulation. Handbook of Fconometrics, 4:2383-2441, 1994.
James J Heckman. Sample selection bias as a specification error. Econometrica: Journal of the Econometric Society, pp. 153-161, 1979.
10

---

Published as a conference paper at ICLR 2023
Dachee Kim, Youngjun Yoo, Seunghyun Park, Jinkyu Kim, and Jackoo Lee. Selfreg: Self- supervised contrastive regularization for domain generalization. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 9619-9628, 2021.
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Bal- subramani, Weihua Hu, Michihiro Yasunaga. Richard Lanas Phillips, Irena Gao, et al. Wilds: A benchmark of in-the-wild distribution shifts. In Proceedings of International Conference on Machine Learning, pp. 5637-5664. PMLR, 2021.
David Krueger, Ethan Caballero, Joern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. Out-of-distribution generalization via risk extrapo- lation (rex). In Proceedings International Conference on Machine Learning, pp. 5815-5826.
PMLR, 2021.
Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy Hospedales. Learning to generalize: Meta-lcarning for domain generalization. In Proceedings of the AAAI Conference on Artificial Intelligence, 2018.
Daniel McFadden. A method of simulated moments for estimation of discrete response models without numerical integration. Econometrica: Journal of the Econometric Society, pp. 995-1026, 1989.
Jose G Moreno-Torres, Troy Raeder, Rocío Alaiz-Rodríguez, Nitesh Chawla, and Francisco Her- rera. A unifying view dataset shift in classification. Pattern Recognition, 45(1):521-530, 2012.
Saeid Motiian, Marco Piccirilli, Donald A Adjeroh, and Gianfranco Doretto. Unified deep su- pervised domain adaptation and generalization. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 5715-5725, 2017.
Krikamol Muandet, David Balduzzi, and Bernhard Schölkopf. Domain generalization via invariant feature representation. In Proceedings of International Conference on Machine Learning, pp- 10-18. PMLR, 2013.
Alexander Robey, George J. Pappas, and Hamed Hassani. Model-based domain generalization. In Proceedings of Advances in Neural Information Processing Systems, pp. 10210-20229 Curran Associates, Inc., 2021.
Elan Rosenfeld, Pradeep Kumar Ravikumar, and Andrej Risteski. The risks of invariant risk mini- mization. In Proceedings of International Conference on Learning Representations, 2021.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenet large scale visual recognition challenge. International Journal of Computer Vision, 115(3))211-255) 2015.
Shiori Sagawa, Pang Wei Koh, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generaliza- tion. In Proceedings of International Conference on Learning Representations, 2020.
Shiori Sagawa, Pang Wei Koh, Tony Lee, Irena Gao, Sang Michael Xie, Kendrick Shen, Ananya Kumar, Weihua Hu, Michihiro Yasunaga, Henrik Marklund, Sara Beery. Etienne David, Ian Stavness, Wei Guo, Jure Leskovec, Kate Saenko, Tatsunori Hashimoto, Sergey Levine, Chelsea Finn, and Percy Liang. Extending the WILDS benchmark for unsupervised adaptation. In Proceedings of International Conference Learning Representations, 2022.
Yuge Shi, Jeffrey Seely, Philip Torr, Siddharth N, Awni Hannun, Nicolas Usunier, and Gabriel Syn- naeve. Gradient matching for domain generalization. In Proceedings of International Conference on Learning Representations, 2022.
Baochen Sun and Kate Saenko. Deep coral: Correlation alignment for deep domain adaptation. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 443-450, Springer, 2016.

---

Published as a conference paper at ICLR 2023
Vladimir Vapnik. Principles of risk minimization for learning theory. In Proceedings of Advances in Neural Information Processing Systems, pp. 831-838. Curran Associates, Inc., 1991.
Chuanlong Xie, Fei Chen, Yue Liu, and Zhenguo Li. Risk variance penalization: From distributional robustness causality. arXiv preprint wwXiv:2006.07544, 2020.
Bianca Zadrozny. Learning and evaluating classifiers under sample selection bias. In Proceedings of International Conference on Machine Learning, pp. 1 114. Association for Computing Machinery, 2004.
Shanshan Zhao, Mingming Gong, Tongliang Liu, Huan Fu, and Dacheng Tao. Domain generaliza- tion via entropy regularization. In Proceedings of Advances in Neural Information Processing Systems, pp. 16096-16107. Curran Associates, Inc., 2020.
Kaiyang Zhou, Yongxin Yang. Timothy Hospedales, and Tao Xiang. Learning to generate novel domains for domain generalization. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 561-578. Springer, 2020.
12

---

Published as a conference paper at ICLR 2023
A APPENDIX
A.I DERIVATION OF EXAMPLE OF DECOMPOSABLE LOss
BU(X), Y) log
1ogp(Y,s*=1115
log |n(Y|X, St = 1)p(5* = =1p(5==1(xx)) =|/x)|}
log =1) logp(s*== = 1|X)}
log Y ( =1) 1)
Σ
log p(s Sh 11x)
A
A.2 PROOF OF THROREM
The following chain of inequalities completes the proof:
1.
E (xxY)~xx
(((SDRM(X)),
E
(X) =1)
(XYY)s
E (x,Y)~x IA(SDRM(XX), gsana (X)D) X);Y,5*)] +
E Exwpx (X) = 1)]P(s*= =1)==1
k=1
E (x
[A(fsana (X), (aa_aaaaax))),,x__ (Xxsnnc(()))xxxxxx))))).x).)
E Exxpx (c)0ss = 1)) =1)
k=1
(13)
(XY))s
k=
E [X, (x
[A(JsanaXX), (X)x)))))) Σ
Exmpy (ooosgll(usaaa(x),,5 = P)) 1)
k=1
[A[JsoooXX) 9SDDG
(14)
(X-Y)~kk
km
E
[ (X)ana(X)) Y)).
(XxY)~Px
Note that 13 follows from the definition of SGDG. 14 holds if Y and S are dependent and if we let
((f(X) Y) logp(Y)XX, A log (Y 1). and lg lllogp(S 1[X).
A.3 DETAILS OF HECKMAN CORRECTION GUIDED DOMAIN GENERALIZATION
In this section, we present the exact form of our loss functions for continuous, binary, and multino- mial outcomes, based on their specific parametric assumptions.
Continuous Outcomes For regression tasks where E I. we assume that the selection and out- come is generated by the following data generation process:
s = I5 > 0 = [[g)(x:) + mh > 0]
(15)
Y f(x) +
(16)
13

---

Published as a conference paper at ICLR 2023
Th
N
0,
opk o²
(17)
apk
where II()) E {0, 1) is the indicator function, Pk is the correlation between 7i and E, and E R. the standard deviation of Being consistent with the notation in the main text, we hereby define {pI....,,P PK, o). Denoting model parameters as = (J.91....55 9K. }), we formulate the data likelihood as follows:
N K
L. D)
= IIII
plu D = 1) p(s; = 1
p(su = =0x,,)' 1-5k
(18)
i= ==
and we take the negative log likelihood to formulate the loss function:
(x,) + Pk Pe f(m=:)
1
{ y/ - f((x)) {²
(<(D:)) =
Sik
log 11
log 2TT
v1
Ph
2
2
σ
8:k) log 11 (x:))
(19)
Binary Outcomes For binary outcomes where y E {0, 1}. we assume a probit model for both selection and outcome:
s* 11$* 0] = [[gh(x) + m > 0]
Y [[f()) +
(20)
Th
N
Pk
Pk 1
where Pk is the correlation between 7}k and Herein, Σ. (pi..... px) Denoting model parame-
ters as 0 = [f.91..... ,K, }), we formulate the data likelihood for binary outcomes as follows:
K (p(: 1, six 1 p(y: sikk 1|xx) 1- si p(sE 01æx,)
Co(0; II
(21)
and take the negative log likelihood define the loss function for binary outcomes:
N K
(D(:0) =
Sik Yi log 22 (gx( f(x;), ps)
K
Sih (1 - 8/) log (41 (gx(xx;) - 4 (ox)
(22)
A=1
K
85k) log +1 (-gi( (xx))
1=1
where P2(, ;0) the cumulative density function of the bivariate standard normal given correlation [-,
Multinomial Outcomes For multinomial outcome tasks where E fl.....J}, we assume multinomial probit outcome model (McFadden, 1989):
S* 11$* 0| = (x)) + mh > 0]
y, = fz) + Ej.
Vj E (1.....J]
Y argmax Y jE(I.....))
(23)
m
px]
1 )
0,
8j
14

---

Published as a conference paper at ICLR 2023
where Ey] (O, Pkyl) and Pkj is the correlation between Ps and Ej- We fur- ther assume that the outcome error terms independently distributed: [1,..,,J] NO,,I). Consequently, {f1,...,
Denoting model parameters as -......... P.k.....PPJJ. ), we formulate the multinomial outcome data likelihood as .......f...1.... fs.91,..... ............ 9K, 9K,,11,.---....,...,,,,,,,,11,,,
(Kk))
follows:
L. (0:
IIIII
II p(sa 1, y: j|xx;) )lvc=)) 8ik p(s )1,=)==>>
(24)
where TT:] E {0, 1 the indicator function. We use the negative logarithm of the data likelihood for the loss function:
N
Sck 1ly: [] log / slu(ū, L)du sa) log 2 (- gx
im (D:0)
(25)
where = |- gb (2;); oo) x (2;), oc) .... ( oo) x [5r+ (æ:). oo) x ... X [5, (a:;), C TR is the half-open J-dimensional hyperrectangular domain and Ejgr = (x:) + fr. We use the GHK algorithm (Hajivassiliou & Ruud, 1994) to compute the mul- tivariate normal integrals.
A.4 OPTIMIZATION, DATASETS, AND MODEL TRAINING
Algorithm Two-Step Optimization for HeckmanDG
1: Input: Data D = { si) E X' x y x S:== N}, Batch size B, Learning rate
2: Output: f.g (9x. ) K 1= and Σ.
3: Initialize: f.g. Σ
4: Step 1: Learn the Domain Selection Models (g)
5: for all 1...., K do
6: DR + {( 8ik) X i=
7: end for
8: while g {g Not Converged do
9:
for all I,... ,.K
K do
10:
B& BatchSampler(De. B)
11:
Is + 9k + 1 V (B, (ss log (gx((x)) + (1 -sw) log ( (x;)))
12:
end for end while
13:
14:
9k = gk, for k I. K.
15:
Step 2: Learn the Outcome Model (J, E)
Do D
17:
while f Not Converged do
Batchsampler
B) VVE 8)) Σ k= 8t. (f (x:); Ve, 88ks E)
19:
f <
20:
Σ t Σ V Σ (E.
E
1= 8a. (f.()) (2::): U.. 882i E)
21:
end while
22:
f+
Σ+Σ x
24:
return
15

---

Published as a conference paper at ICLR 2023
Dataset
CAMELYON1?
POVERTYMAP
IWILDCAM
RXRXI
X
3 x 96 x 96 (lisuuc slide)
8 x 224 x 224 (satelliie image)
448 x 448 (photo)
3 x 256 x 256 (well)
2 (Numor)
contimuons (avset wealth)
182 (animal species) 129,809
1139 (genetic treatments)
Triining examples
502,436
10,000
40,612
Domain
Hospitals (3, 1)
23 countries (13,
camarai inaps (243,
51 batches (33, 4, 14) Averago accuneyy
Evaltmtion metric
Average acearney
Peansio (everage worst-grupp
Maero
Table 5: A summary on the four datasets from WILDS benchmark. In the 'Domain' row, the three numbers in parentheses denote the number of train, validation, and test domains.
Dataset
CAMELYON17
POVERTYMAP
IWILDCAM
RxRxI
Feature Extractor
DenseNet-121
ResNet-18-MM
ResNet-50
ResNet-50
Epochs
5.5
100, 100
12, 12
10, 10
Batch Size
32, 32
64, 64
16, 16
75, 75 Adam, SGD
Optimizer
Adam, SGD
Adam, SGD
Adam, SGD
Learning Rate
10-9. 10-3
10 10
10 10-0, 10-3 0, 10 True, True
10 - 10-3 0, 10 True, True RandAugment
Weight Decay
0, 10
0, 10
ImageNet Weights
True, True N/A
False, False
Data Augmentation
Color jittering
RandAugment
Table 6: Details on training configurations of HeckmanDG. Cells with two entries (i.e., learning rate, weight decay) denote that used different values for training domain selection and outcome mod- els. Hyperparameters were determined through grid search on the OOD validation set. For optimiz- ers we searched between (Adam, SGD}. Learning rates were searched among {10 5 10-4 10 3) and weight decay among (0, 10 10-5, 10-3 110-1,, Note that since HeckmanDDG training is two-phase, we searched for hyperparameters sequentially. On all datasets, optimized with the Adam opti- mizer with learning rate of 10-2 and no weight decay.
A.5 IMPLEMENTATION DETAILS FOR SWAD AND L2A-OT
In Tables 2, 3, and 4, the numbers for SWAD (Cha et al., 2021) and 1.2AOOT (Zhou et al., 2020) reproduced based on implementation. For SWAD, set the optimum patient parameter N. 3, 3, the overfitting patient parameter N. 6, and the tolerance rate = 1.2, which are the default values used in the original paper. The evaluation frequency was set 100. On all four datasets, models were trained with SGD, using the same learning rates and weight decay factors as ERM (Koh et al., 2021). For L2A-OT, we set Apomain = 0.5, Acyele = 10, and ACE 1. On all three datasets (excluding POVERTYMAP), the generator G is trained with Adam using a constant learning rate of 3- 10-4, while the prediction model F trained with SGD, using the same learning rates and weight decay factors as ERM (Koh et al., 2021). We faithfully refer the readers to the original papers for the details regarding the hyperparameters.
16

---

Published as a conference paper at ICLR 2023
A.6 ADDITIONAL ANALYSIS ON TWO-STEP OPTIMIZATION
We provide further details on the analysis in Section 7, conducted to demonstrate the necessity and effectiveness of the proposed two-step optimization approach. For both one-step (jointly optimizing for Equation 8) and two-step, we kept the number of training epochs to 30, with the same batch size, order to make valid comparison. In Figure 4, we plot the training loss trajectories using different metrics. We observed that tends to converge to a suboptimal point if we optimize f. g, and E simultaneously by one-step optimization approach (Figure 4a). We suspect that this happens because the gradient of (involving $ = 1) depends on Σ, which changes during the joint one- step optimization process. This will often result suboptimal g, as shown in Figure 4d. Note that the original Heckman correction (Heckman, 1979) also proposed two-step approach to avoid the computational burden of having to estimate both f and jointly. One way to avoid this problem is to use alternating optimization algorithm which alternately minimizes Equation 8 with respect to g, Σ, and f until convergence. Our two-step approach can be regarded as a one-iteration alternating minimization procedure it stops after single iteration. In Figure 4c, the one-step method yields comparable (MSE) loss values for f as those obtained from the two-step method on the training dataset. However Equation converged to larger value by the one-step method, implying that converged to a suboptimal point for g. On the test datasets, the suboptimal solution of the one- step method resulted in worse domain generalization performance compared the two-step method (Figure 3b).
ore step
3.0
on stop
two step
3mo stop
2.5
Tran lowe
0.3
Tran
2.0
1.5
1.0
(a) Train loss (Equation 8)
(b) Train A
stop
stop
Thiin MAVE 09
0.8
0.3
0.6
0.2
0.4
6.1
0.0
10
15
20
25
30
10
25
Epechs
Epochs
(c) Train MSE
(d) Train loss of g (Equation 11)
Figure 4: Training loss of one-step (gray) vs. two-step (blue) optimization on POVERTYMAP.
17


Reviews:
[{"review_id": "sdVDvsAnTdP", "reviewer": "Reviewer_Y91A", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "The authors consider domain generalization as a non-random sample selection problem. It is the first paper of kind to utilize this method for DG problem. Via assessment on both simulation and benchmarking data sets, the authors demonstrate the efficacy of our method both theoretically and empirically on simulated data and four challenging benchmarks.", "summary_of_the_review": "The research work is well conducted and the paper is clearly written. Besides a couple of weakness, the research results are worth sharing with the broader audiences to raise discussions and potential further research in this direction.", "strength_and_weaknesses": "The paper proposed a new method and did extensive experiments to demonstrate the effectiveness of the proposed method, which is a strength of the paper.\n\nAs the authors stated at the beginning of the paper, DG questions may be categorized into different categories. While the paper shows that the proposed method works in many cases, it is unclear to the reader, what exact scenario the proposed method may work and in which scenarios it will not. If the authors can shed some light on this direction, it will better navigate the readers about future applications of the method.\n\nIt will also be good if the authors can describe the pre-assumption of the methods more explicitly. The original Hackman's method has a number of assumptions, including the assumption of normal distribution etc. What are the implications for adoptions in the current scenarios? Hope the authors can add some discussions on such aspects.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "r5xGS7pl7cn", "reviewer": "Reviewer_CmDk", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "This paper proposes a method for domain generalization in which the variable $S$ selecting the domain correlates both the input $X$ and the output variable $Y$. The proposed method models the joint probability distribution of $(X, Y, S)$ and performs a two-step variant of the maximum likelihood estimation (MLE). The authors demonstrate that the proposed method outperforms the ERM as well as other domain generalization methods. Interestingly, the proposed two-step algorithm performs better than jointly performing the MLE.", "summary_of_the_review": "The proposed method is interesting and seemingly effective, but it looks like there is a gap between the theory part and the proposed algorithm. I weakly recommend acceptance.", "strength_and_weaknesses": "# Strengths\n- The approach of exploiting the correlation between the data and the domain is interesting.\n- Theorem 1 is also interesting and well supports the proposed approach.\n- The experiments show that the proposed method significantly outperforms previous approaches.\n\n# Weaknesses\n- In Section 4, \"Equation 5 becomes...\": I think this is related to the previous point. Eq. (7) does not exactly correspond to Eq. (5) because it ignores $k = K + 1, \\dots, L$. But if we ignore those terms, we will go back to Eq. (3). I am not sure if the proposed method is doing what the paper claims it does.\n- The two-step algorithm does not exactly solve Eq. (8) but performs better than the more theoretically faithful one-step algorithm, according to Section 7. The paper does not provide a convincing explanation about this. I suspect that the two-step algorithm is solving something different from Eq. (8) but theoretically justifiable.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "ehLFqYJ_VA", "reviewer": "Reviewer_2hmL", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "This paper casts domain generalization as a sample selection problem. It starts from problem formulation which is then disentangled into two terms: sample selection loss and joint function loss. Authors designed a dual-branch network to implement this idea. Experiments are done on WILDs benchmark, showing its effectiveness against its counterparts.", "summary_of_the_review": "I think this paper is novel in formulating the DG problem as a sample selection one. It has its own theory and derivation, with working algorithms. \n\n=== Comments after rebuttal: I decided to increase the score to 8 given that authors addressed my concerns well.", "strength_and_weaknesses": "### Strength\n\n1. This paper is well motivated. It is interesting to see that authors formulate this problem as a sample selection problem.\n2. The theory part is novel, with solid proof.\n3. Experiments show its effectiveness.\n\n### Weakness\n\n1. The computation complexity is not introduced in the paper. From what I see, there are two feature extractors in the method, while comparison methods only adopt one. So, what's the time complexity and number of parameters comparison? Are the experiments fair?\n2. The latest comparison method is Fish, which is in 2021. There are other approaches more latest than this: SWAD and L2A-OT, for instance. Authors should cite and compare with them\n3. There are no ablation studies to show the actual contribution of each component in the method.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%