Original: D OMAIN G ENERALIZATION VIA H ECKMAN - TYPE S ELECTION M ODELS
Revised: D OMAIN G ENERALIZATION VIA H ECKMAN - TYPE S ELECTION M ODELS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: The proposed SGDG is domain generalizable as it intends to minimize the risk under the population distribution.
Revised: The proposed SGDG is domain generalizable as it intends to minimize the risk under the population distribution.

Original: We theoretically proved that, under certain regular conditions, SGDG can achieve smaller risk than
Revised: We theoretically prove that, under certain regular conditions, SGDG can achieve smaller risk than

Original: ERM.
Revised: ERM.

Original: Furthermore, we present a class of parametric SGDG (HeckmanDG) estimators applicable to continuous, binary, and multinomial outcomes.
Revised: Furthermore, we present a class of parametric SGDG (HeckmanDG) estimators applicable to continuous, binary, and multinomial outcomes.

Original: We also demonstrated its efficacy empirically through simulations and experiments on a set of benchmark datasets comparing with other well-known DG methods.
Revised: We also demonstrate its efficacy empirically through simulations and experiments on a set of benchmark datasets comparing with other well-known DG methods.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: 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 same underlying data distribution (Vapnik, 1991).
Revised: 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 same underlying data distribution (Vapnik, 1991).

Original: In other words, most statistical models assume that the training and test data are both random samples chosen randomly from the same population.
Revised: In other words, most statistical models assume that the training and test data are both random samples chosen randomly from the same population.

Original: Unfortunately, this assumption is often violated in real-world applications rendering model performance to deteriorate on out-of-distribution (OOD) test data (Koh et al., 2021).
Revised: Unfortunately, this assumption is often violated in real-world applications rendering model performance to deteriorate on out-of-distribution (OOD) test data (Koh et al., 2021).

Original: Recently, the
Revised: Recently, the

Original: Domain Generalization (DG) problem (Blanchard et al., 2011) 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 to new instances sampled from previously unseen test domains.
Revised: Domain Generalization (DG) problem (Blanchard et al., 2011) 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 to new instances sampled from previously unseen test domains.

Original: For instance, it is hard to expect a learning algorithm for predicting whether a patient is diagnosed with cancer, trained on data obtained from thelocal hospitals in region A via a specific data collection protocol, to perform equally well on data collected at hospitals in region B, using different data-collection techniques. In this context, cancer patients are the underlying population of interest.
Revised: In this paper, we consider DG as a non-random sample selection problem. Let P XY represent the population data distribution , and S k denote a binary random variable indicating whether a subject is selected from the population into domain k .

Original: Both the patient samples in hospital A and B are selected non-randomly from the population, and their selection probabilities most likely differ.
Revised: In a random sampling process, P ( S ki = 1) is independent and identically distributed ( iid ).

Original: In this paper, we consider domain generalization as a non-random sample selection problem. The samples ineach domain are not selected randomly from the true population, but rather under a selection probability, P (  S  = 1) , that is correlated with both the predictive features X and the outcome Y . 
Revised: Under a non-random sample selection, the distribution of ( X, Y ) in domain k is a conditional distribution of ( X, Y ) given S k = 1 , which often does not equal to P XY . Consequently, this leads to distributional shifts across domains: P jXY ̸ = P kXY for k ̸

Original: A naively trained model, optimizing the prediction performance conditional on the selection mechanisms of the domains in the training data, will fail to perform similarly if sample selection of a testing domain is under a different selection mechanism.
Revised: = j .

Original: In the statistical and machine learning literature, methods for DG often build on different assumptions regarding the source of distributional shift. The vast majority of methods rely on the covariate assumption (Bickel et al., 2009; Castro et al., 2020), where domains differ only in their marginal data distributions  P  (  X ) .
Revised: Mathematically, distribution shifts across domains P kXY may contain shifts in distributions of X ( P kX , covariate shift (Bickel et al., 2009), and in the distributions of Y conditional on X ( P kY | X , concept shift (Moreno-Torres et al., 2012). We present graphical models in Figure 1 to conceptually illustrate the sources of distribution shifts, assuming the existence of latent factors confounding the relationship between X , Y , and domain ( S k ). In Figure 1, C 1 represents unobserved latent factors that correlate with X and S , resulting in covariate sfhift.

Original: Some make the concept shift assumption, where the conditional distribu- tion P ( Y | X ) differs across domains.
Revised: C 2 correlates with X , Y , and S simultaneously, entailing both covariate and concept shifts. The goal of DG is to estimate the domain generalizable (agnostic) edge f : X → Y in the presence of the two types of latent confounders.

Original: Perhaps a more interesting and understudied case would be when we can make no assumptions on independence between Y and S (Koh et al., 2021), as the outcome is  often correlated with sample selection.
Revised: The vast majority of DG methods are developed to identify f that is robust to C 1 . However in practice, C 2 type of confounders often exist which make P ( S k = 1) related to both X and Y . For [Figure]

Original: Most notably, Heckman (Heckman, 1979) proposed a bias correction approach to formulate the joint distribution of the outcome and the selection probability P (Y, S | X ) , then use this probability to learn an unbiased model of the true population.
Revised: example, when we train a model to predict tumor status ( Y ) from histological images ( X ) using patients from different hospitals ( S k ), there may be variations in X due to inconsistent acquisition processes such as staining differences ( C 1 ) across hospitals, and differences in patient characteristics such as age, gender, race, and disease severity ( C 2 ) that correlate with hospital, covariates and the outcome. As a result, a model trained in an oncology specialist hospital may not be generalizable to a hospital serving veterans. Similarly, when we train a 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 ( C 2 ) correlating with X, Y and domains simultaneously. Therefore, a model trained on one country may not perform well on another country with a different rural/urban proportion or economic status (Koh et al., 2021).

Original: In this paper, we proposed a new class of Selection Guided Domain Generalization (SGDG) models to first estimate the  probability that an instance is sampled into a training domain, and then use the joint distribution of the outcome Y andsample selection S to learn a domain generalizable model.
Revised: 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.

Revised: In particular, SGDG is built on Heckman’s bias correction framework (Heckman, 1979) which is a very powerful tool to learn an unbiased model from non-randomly selected samples in the presence of both C 1 and C 2 confounders.

Original: The unique contributions of this paper are summarized as follows:
Revised: The unique contributions of this paper are summarized as follows:

Original: • 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.
Revised: • 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.

Original: • We present a class of parametric SGDG (HeckmanDG) estimatorsto jointly model the selection and outcome mechanisms, applicable to continuous, binary, and multinomial outcomes.
Revised: • We present a class of parametric SGDG (HeckmanDG) estimators applicable to continuous, binary, and multinomial outcomes.

Original: • We demonstrate the efficacy of our method both theoretically and empirically on simulated data and four challenging benchmarks.
Revised: • We demonstrate the efficacy of our method both theoretically and empirically on simulated data and four challenging benchmarks.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Domain Generalization.
Revised: Domain Generalization.

Original: DG has been studied under various contexts, where each proposing sophisticated learning objectives that leverage domain information to bypass spurious correlations (Sun & Saenko, 2016; Sagawa et al., 2019; Arjovsky et al., 2019; Shi et al., 2021; Arpit et al., 2021).
Revised: DG has been studied under various contexts.

Original: A considerable body of research has been devoted to learning domain-invariant features which are discriminative of the task but at the same time independent of the domain. Such methods include but are not limited to 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 et al., 2020), and contrastive learning (Motiian et al., 2017; Kim et al., 2021).
Revised: Many studies are devoted to learning domain-invariant 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 et al., 2020), and contrastive learning (Motiian et al., 2017; Kim et al., 2021).

Original: Meanwhile, worksinspired by causality exploit invariant causal effects across domainsto avoid learning spurious correlations via learning feature representations on top of which a single optimal predictor exists (Arjovsky et al., 2019; Ahuja et al., 2020; Rosenfeld et al., 2020).
Revised: Other works exploit invariant causal effects across domains (Arjovsky et al., 2019; Ahuja et al., 2020; Rosenfeld et al., 2020).

Original: Another family of robust optimization methods seek toiteratively minimize the worst-case errorover an uncertainty set of data distributions constructed from the training domains (Sagawa et al., 2019; Xie et al., 2020; Krueger et al., 2021).
Revised: Another family of robust optimization methods seek to minimize the worst-case error (Sagawa et al., 2019; Xie et al., 2020; Krueger et al., 2021).

Original: More recently, another prominent direction of methods focus on model averaging throughout single or multiple training trajectories (Cha et al., 2021; Arpit et al., 2021).
Revised: More recently, other prominent directions of methods improve

Original: Other approaches include gradient matching (Shi et al., 2021) and meta learning-based methods (Li et al., 2018). The computer vision community has also presented a series of methods, especially based on data augmentation (Robey et al., 2021).
Revised: DG by model averaging (Cha et al., 2021; Arpit et al., 2021), gradient matching (Shi et al., 2021), meta learning (Li et al., 2018), data augmentation (Robey et al., 2021), and generating novel domains (Zhou et al., 2020).

Original: Sample Selection Bias Correction.
Revised: Sample Selection Bias Correction.

Original: Zadrozny (2004) formalized sample selection bias in machine learning terms and presented a bias correction method for the case where selection only depends on the input features.
Revised: Zadrozny (2004) formalized sample selection bias in machine learning terms and presented a bias correction method when selection only depends on the input features.

Original: Cortes et al.
Revised: Cortes et al.

Original: (2008) proposed a sample reweighting approach to tackle the same problem, but assumes theadditional availability of unlabeled data drawn  the true population.
Revised: (2008) proposed a sample reweighting approach to tackle the same problem, but assumes the availability of additional data drawn from the true population.

Original: Cortes & Mohri (2014) tackles domain adaptation and sample bias correction in regression.

Original: Sample reweighting approaches for robust prediction have also been considered, but not under the problem of domain generalization (Liu & Ziebart, 2014).

Original: Du & Wu (2021) proposed a framework for robust and fair learning under  sample selection bias, but assumes conditional independence of Y and S given X .
Revised: Du & Wu (2021) proposed a framework for robust and fair learning under biased sample selection, but assumes conditional independence of Y and S given X .

Original: However, it is often the case that in real world domain generalization problems, the outcome is correlated with sample selection.

Original: To this end, our work considers the case where there is no conditional independence relationships between X , Y , and S .

Original:  S ELECTION M ODEL -G UIDED D OMAIN G ENERALIZATION 3. D OMAIN G ENERALIZATION
Revised:  S ELECTION M ODEL -G UIDED D OMAIN G ENERALIZATION 3. D OMAIN G ENERALIZATION

Original: Suppose there exists L distinct but relevant domains and let S
Revised: Suppose there exists L distinct but relevant domains and let S

Original: =
Revised: =

Original: ( S 1 , · · · , S L ) ∼ P S denote a binary random vector that indicates the domain membership where S k = 1 implies belonging to domain k .
Revised: ( S 1 , · · · , S L ) ∼ P S denote a binary random vector that indicates the domain membership where S k = 1 implies belonging to domain k .

Original: Let P XY represent the population data distribution , and each domain’s data distribution be a conditional distribution of the population distribution given S k = 1 .
Revised: Let P XY represent the population data distribution , and each domain’s data distribution be a conditional distribution of the population distribution given S k = 1 .

Original: i.e., P kXY =
Revised: i.e., P kXY =

Original: P XY | S k
Revised: P XY | S k

Original: =1 .
Revised: =1 .

Revised: Assumption 1 (Mutually Exclusive Domain Membership) .

Original: We assume that if S k = 1  then S j = 0 for all j ̸ = k
Revised: We assume that if S k = 1 , then S j =for all j ̸ = k

Original: so that an example can be observed in one and only one domain.
Revised: so that an instance can belong to one and only one domain.

Revised: Assumption 2 (Independent Domain Sampling Processes) .

Revised: We assume that S k ⊥⊥ S j for all j ̸ = k .

Original: In the supervised domain generalization context, we are allowed to observe the joint distribution of
Revised: In the supervised domain generalization context, we are allowed to observe the joint distribution of

Original: X and Y , P kXY for K out of L domains, and refer the K domains as source or training domains observed during the training phase.
Revised: X and Y , P kXY for K out of L domains, and refer the K domains as source or training domains observed during the training phase.

Original: The remaining L − K domains are referred to as target or testing domains whom we may observe in the testing phase.
Revised: The remaining L − K domains are referred to as target or testing domains whom we may observe in the testing phase.

Original: Under this setting, we aim to learn a prediction model that generally performs well on both source and target domains.
Revised: Under this setting, we aim to learn a prediction model that generally performs well on both source and target domains.

Original: We formalize the problem as follows.
Revised: We formalize the problem as follows.

Original: Definition 1 (Domain Generalization (Blanchard et al., 2011)) .
Revised: Definition 1 (Domain Generalization (Blanchard et al., 2011)) .

Original: 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:
Revised: 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:

Original: [Equation]
Revised: [Equation]

Original: where ℓ : Y × Y → R + (= { r ∈ R : r ≥ 0 } ) is a loss function and F is a hypothesis set.
Revised: where ℓ : Y × Y → R + (= { r ∈ R : r ≥ 0 } ) is a loss function and F is a hypothesis set.

Original: We note that some other papers have considered the same problem setting (Muandet et al., 2013;
Revised: We note that some other papers have considered the same problem setting (Muandet et al., 2013;

Original: Deshmukh et al., 2019; Blanchard et al., 2021).
Revised: Deshmukh et al., 2019; Blanchard et al., 2021).

Original: We first introduce a proposition that claims the equivalence of the domain generalization and the risk minimization under the population distribution:
Revised: We first introduce a proposition that claims the equivalence of the domain generalization and the risk minimization under the population distribution:

Original: Proposition 1 (Equivalence of Domain Generalization and Population Risk Minimization) .
Revised: Proposition 1 (Equivalence of Domain Generalization and Population Risk Minimization) .

Original: Problem (1) is equivalent to the risk minimization under the population distribution P XY .
Revised: Problem (1) is equivalent to the risk minimization under the population distribution P XY .

Original: That is,
Revised: That is,

Original: [Equation]
Revised: [Equation]

Original: which straightforwardly follows from the law of total expectation.
Revised: which straightforwardly follows from the law of total expectation.

Original: We define f PRM as the minimizer of the population risk minimization problem, as well as the best (hypothetical) model for the domain generalization problem.
Revised: We define f PRM as the minimizer of the population risk minimization problem, as well as the best (hypothetical) model for the domain generalization problem.

Original: Proposition 1 is important as it establishes that the best model for the domain generalization problem should minimize the risk under the population distribution P XY .
Revised: Proposition 1 is important as it establishes that the best model for the domain generalization problem should minimize the risk under the population distribution P XY .

Original: However, in practice we have no access to the population data distribution and are only given the source domains during model training.
Revised: However, in practice we have no access to the population data distribution and are only given the source domains during model training.

Original: Therefore, one naive learning problem is minimizing the risk under the source domain data distribution:
Revised: A naive learning model is minimizing the risk under the source domain data distribution:

Original: Definition 2 (Source Domain Risk Minimization) .
Revised: Definition 2 (Source Domain Risk Minimization) .

Original: Given the source domains whose data distributions are P kXY =
Revised: Given the source domains whose data distributions are P kXY =

Original: P XY | S k
Revised: P XY | S k

Original: =1 , for k = 1 , · · · , K , we refer to the following learning problem as the source domain risk minimization problem
Revised: =1 , for k = 1 , · · · , K , we refer to the following learning problem as the source domain risk minimization problem

Original: [Equation]
Revised: [Equation]

Original: We denote its minimizer as f SDRM .
Revised: We denote its minimizer as f SDRM .

Original: The empirical form of f SDRM is denoted as empirical risk minimization (ERM).
Revised: The empirical form of f SDRM is denoted as empirical risk minimization (ERM).

Original: The generalization performance of f SDRM depends on how well the source domains represent the population (or the target domains).
Revised: The generalization performance of f SDRM depends on how well the source domains represent the population (or the target domains).

Original: If the source domains well approximate the population distribution, then the generalization performance of f SDRM will be sufficiently close to that of f PRM .
Revised: If the source domains well approximate the population distribution, then the generalization performance of f SDRM will be sufficiently close to that of f PRM .

Original: On the other hand, if P kXY ̸ = P XY , one cannot guarantee that models trained on the source domains to effectively generalize to the population.
Revised: On the other hand, if P kXY ̸ = P XY , one cannot guarantee that models trained on the source domains to effectively generalize to the population.

Original: Therefore, it is necessary to model the selection probability to bridge the gap between P kXY and P XY .
Revised: Therefore, it is necessary to model the selection probability to bridge the gap between P k XY and P XY .

Original: 3. S ELECTION M ODEL -G UIDED D OMAIN G ENERALIZATION
Revised: 3. S ELECTION M ODEL -G UIDED D OMAIN G ENERALIZATION

Original: Now we derive our selection model-guided domain generalization problem, starting from decomposing the objective function of the domain generalization problem by introducing a domain selection model
Revised: Now we derive our selection model-guided domain generalization problem, starting from decomposing the objective function of the domain generalization problem by introducing a domain selection model

Original: g = {
Revised: g = {

Original: g k : X → [0 , 1] , k = 1 , · · · , L } which predicts selection probabilities of an example being observed in domain k .
Revised: g k : X → [0 , 1] , k = 1 , · · · , L } which predicts selection probabilities of an example being observed in domain k .

Revised: Assumption 3 (Decomposable Loss Function) .

Original: We assume that the loss function ℓ can be decomposed into two components - one exclusively about the selection model and the other involving both models:
Revised: We assume that the loss function ℓ can be decomposed into two components - one exclusively about the selection model and the other involving both: ℓ

Original: [Equation]
Revised: [Equation]

Original: where ℓ s : [0 , 1] L ×{ 0 , 1 } L → R + is the loss function for learning the domain selection models and Λ : ( Y × [0 , 1] L )
Revised: where ℓ s : [0 , 1] L ×{ 0 , 1 } L → R + is the loss function for learning the domain selection models and Λ : ( Y × [0 , 1] L )

Original: ×
Revised: ×

Original: ( Y × { 0 , 1 } L ) → R + is the joint loss function, which assigns to the prediction and domain selection models a pair of true outcome Y and domain membership indicator S .
Revised: ( Y × { 0 , 1 } L ) → R + is the joint loss function, which assigns to the prediction and domain selection models a pair of true outcome Y and domain membership indicator S .

Original: For example,  under the probabilistic framework, if we let ℓ ( f ( X ) , Y ) =
Revised: For example, the negative log-likelihood under the probabilistic framework satisfies Assumption 3: if we let ℓ ( f ( X ) , Y ) =

Original: − (cid:80)
Revised: − (cid:80)

Original: Lk
Revised: Lk

Original: =1 { log p
Revised: =1 { log p

Original: ( Y | X, S k
Revised: ( Y | X, S k

Original: = 1) + log p ( S k = 1 | X ) } , then Λ = − log p
Revised: = 1) + log p ( S k = 1 | X ) } , then Λ = − log p

Original: ( Y | X, S k = 1) and ℓ s = − log p ( S k = 1 | X ) , respectively.
Revised: ( Y | X, S k = 1) and ℓ s = − log p ( S k = 1 | X ) , respectively.

Original: We present a specific parametric form in Section 4.
Revised: We present a specific parametric form in Section 4.

Original: Under the decomposition, the population risk can be expanded as follows:
Revised: Under Assumption 3, the population risk can be expanded as follows:

Original: [Equation]
Revised: [Equation]

Original: Based on the expansion, we introduce our selection model-guided domain generalization problem.
Revised: Based on the expansion, we introduce our selection model-guided domain generalization problem.

Original: Definition 3 (Selection Model-Guided Domain Generalization (SGDG)) .
Revised: Definition 3 (Selection Model-Guided Domain Generalization (SGDG)) .

Original: Given access to the data distribution of K source domains, P kXY for k = 1 , · · · , K and to the unlabeled data distribution from all L domains, P X | S k
Revised: Given access to the data distribution of K source domains, P kXY for k = 1 , · · · , K and to the unlabeled data distribution from all L domains, P X | S k

Original: =1 for k = 1 , · · · , L .
Revised: =1 for k = 1 , · · · , L .

Original: We define the selection model-guided domain generalization problem as a joint learning problem of f : X → Y and g = {
Revised: We define the selection model-guided domain generalization problem as a joint learning problem of f : X → Y and g = {

Original: g k } Lk
Revised: g k } Lk

Original: =1 , min
Revised: =1 , min

Original: [Equation]
Revised: [Equation]

Revised: ∈F , g ∈G

Revised: K

Original: [Equation]
Revised: [Equation]

Original: where F and G are hypothesis sets.
Revised: where F and G are hypothesis sets.

Original: 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 g under the joint distribution of X and S .
Revised: 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 g under the joint distribution of X and S .

Original: Under this formulation, the domain selection model g guides f to be corrected through Λ , considering the probability of being drawn from certain domains.
Revised: Under this formulation, the domain selection model g guides f to be corrected through Λ , considering the probability of being drawn from certain domains.

Original: In Section 4, we will introduce specific forms of this problem under parametric assumptions.
Revised: In Section 4, we will introduce specific forms of this problem under parametric assumptions.

Original: We note that SGDG partially minimizes the population risk (4) ignoring its last term.

Original: It is because, in DG settings, we have no access to the distribution of Y given k = K + 1 , · · · , L .

Original: Under the assumption that the true outcome model f : X → Y is the same for the source domains k = 1 , · · · , K , to the unobserved domains k = K + 1 , · · · , L , f SGDG minimizes the population risk.

Original: Even if the assumption does not hold, we can still prove the following theorem.

Original: Theorem 1 (Performance Improvement of SGDG over SDRM) .
Revised: Theorem 1 (Performance Improvement of SGDG over SDRM) .

Original: Let f SDRM and f
Revised: Let f SDRM and f

Original: SGDG be defined as in Definition 2 and 3.
Revised: SGDG be defined as in Definition 2 and 3.

Original: f SGDG has lower risk than that of f SDRM .
Revised: f SGDG has lower risk than that of f SDRM .

Original: That is,
Revised: That is,

Original: [Equation]
Revised: [Equation]

Original: which implies SGDG is expected to show better generalization performance compared to SDRM.
Revised: which implies SGDG is expected to show better generalization performance compared to SDRM.

Original: By Proposition 1, SGDG is expected to perform better than SDRM for the domain generalization problem equivalently.
Revised: By Proposition 1, SGDG is expected to perform better than SDRM for the domain generalization problem equivalently.

Original: The sketch of proof of this theorem is that SGDG performs as well as SDRM on the source domains and offers performance improvement on the target domains.
Revised: The sketch of proof of this theorem is that SGDG performs as well as SDRM on the source domains and offers performance improvement on the target domains.

Original: Proof of this theorem is given in Appendix A.1.
Revised: Proof of this theorem is given in Appendix A.1.

Original: In reality, we may have no access to the distribution P kX for all domains ( k = 1 , · · · , L ) but only have access to the source domains’ distributions ( k = 1 , · · · , K ). 
Revised: In reality, we may have no access to the distribution P kX for all domains ( k = 1 , · · · , L ) but only have access to the source domains’ distributions ( k = 1 , · · · , K ). In such cases, we use P kX for k = 1 , · · · , K to learn the selection model g as in Problem (7). Assumption 4.

Original: In such cases, we use P kX for k = 1 , · · · ,  K as a surrogate, implying that we learn selection models by contrasting one domain to the remaining source domains.
Revised: Let g ∗ be the optimal selection model that has the minimum expected ℓ s , [Equation] We assume that the first K coordinates ( g ∗ 1 , · · · , g ∗ K ) minimize (cid:80) Kk [Equation]

Original: We empirically confirmed that even with this surrogate, SGDG outperforms the SDRM approach.
Revised: 1 E X ∼P kX [ [Equation] s ( g ( X ) , S k = 1)] P ( S k = 1) .

Revised: Conceptually, Assumption 4 means that the same selection models g k can be learned by contrasting domain k to the remaining domains in the training data ( K \ { k } ) as by contrasting domain k to the remaining domains in the population ( L \ { k } ).

Revised: Under Assumption 4, Problem (5) reduces to min

Revised: [Equation]

Original:  H ECKMAN -T YPE S ELECTION -G UIDED D OMAIN G ENERALIZATION
Revised:  H ECKMAN -T YPE S ELECTION -G UIDED D OMAIN G ENERALIZATION

Original: In this chapter, we present parametric models for Λ embodying f SGDG presented in the previous section.
Revised: In this chapter, we present parametric models for Λ embodying f SGDG presented in the previous section.

Original: The essence of Λ is to model the conditional distribution of P ( Y
Revised: The essence of Λ is to model the conditional distribution of P ( Y

Original: | S
Revised: | S

Original: = 1 , X ) .
Revised: = 1 , X ) .

Original: Consider the setting where we are given training data from multiple source domains in the form of D = { x i , s i , y i } Ni =1 , where s i =
Revised: Consider the setting where we are given training data from multiple source domains in the form of D = { x i , s i , y i } Ni =1 , where s i =

Original: [ s i 1 , s i 2 , . . .
Revised: [ s i 1 , s i 2 , . . .

Original: , s ik , . . .
Revised: , s ik , . . .

Original: , s iK ] ∈ { 0 , 1 } K is a binary vector indicating domain membership.
Revised: , s iK ] ∈ { 0 , 1 } K is a binary vector indicating domain membership.

Original: By formulating the loss function as the negative log-likelihood, the empirical form of Equation 5 becomes
Revised: By formulating the loss function as the negative log-likelihood, the empirical form of Equation 7 becomes

Original: [Equation]
Revised: [Equation]

Original: Heckman (1979) proposed to model the joint distribution P ( Y, ˜ S | X ) of the selection latent variable ˜ S ,
Revised: Heckman (1979) proposed to model the joint distribution P ( Y, ˜ S | X ) of the selection latent variable ˜ S ,

Original: where S =
Revised: where S =

Original: I
Revised: I

Original: [ ˜ S > 0] , and the continuous outcome Y via a bivariate normal distribution with a correlation coefficient ρ  .
Revised: [ ˜ S > 0] , and the continuous outcome Y via a bivariate normal distribution with a correlation coefficient ρ and the mean as a linear functions of X .

Original: The means of the bivariate normal distributionare assumed to be linear functions of X .
Revised: Building upon his work, we make the following assumption on the joint distribution of the outcome and the selection variables. Assumption 5 (Joint Distribution of Y (or latent variables ˜ Y ) and ˜ S ) .

Original: Building upon his work, we use the latent continuous variable  ˜ S k  where S k=
Revised: Let Y be the outcome variable and ˜ S k be the latent continuous variable where S k

Original:  I [ ˜ S k > 0] for each domain k .
Revised: = I [ ˜ S k > 0] for each domain k .

Original: We then model the joint distribution P (Y, ˜ S k | X ) of the selection latent variables and the continuous outcome Y via a multivariate normal distribution with a series of correlation coefficients { ρ k } Kk =1 . The means of the multivariate normal distribution are assumed to be f ( X ) and g k ( X) .
Revised: We assume that Y and S = ( S 1 , · · · , S K ) are jointly distributed as a multivariate normal distribution with mean ( f ( X ) , g k ( X )) , and correlation coefficients { ρ k } Kk =1 , given X .

Original: We assume that { g k } Kk =1 are a set of independent probit models.This can be extended to binary and multinomial outcomes by assuming latent continuous variables underlying the observed outcomes.
Revised: For binary and multinomial outcomes, this assumption is on ˜ Y as the latent continuous variables underlying the observed outcomes. Under this assumption, g = { g k } Kk =1 can be modeled as a set of independent probit models.

Original: In all three cases, the joint log-likelihoods have closed forms.
Revised: In all three cases, the joint log-likelihoods have closed forms.

Original: Henceforth, we refer to the specific parametric form of f SGDG as the Heckman-type DG (HeckmanDG) estimator.
Revised: Henceforth, we refer to the specific parametric form of f SGDG as the Heckman-type DG (HeckmanDG) estimator.

Original: Definition 4 (Heckman-Type Domain Generalization Estimator) .
Revised: Definition 4 (Heckman-Type Domain Generalization Estimator) .

Original: We formulate HeckmanDG as a joint learning problem of f and g = { g k } Kk =1 with the following learning objective:[Equation]
Revised: We formulate HeckmanDG as a joint learning problem of f and g = { g k } Kk =1 with the following learning objective:

Original: [Equation]
Revised: [Equation]

Original: where Φ( · ) is the cumulative distribution function of the standard normal distribution ϕ , such that Φ( g k ( x i ))
Revised: where Φ( · ) is the cumulative distribution function of the standard normal distribution ϕ , such that Φ( g k ( x i ))

Original: =
Revised: =

Original: P ( S k = 1 | X = x i ) is the selection probability w.r.t domain k , and Φ( − g k ( x i )) =
Revised: P ( S k = 1 | X = x i ) is the selection probability w.r.t domain k , and Φ( − g k ( x i )) =

Original: P ( S k = 0 | X = x i ) .
Revised: P ( S k = 0 | X = x i ) .

Original: Meanwhile, Λ( f ( x i ) , g k ( x i ) , y i , s ik ; Σ) is the conditional negative log probability of y i given s ik = 1 , i.e., − log p ( y i | s ik = 1 , x i ) .
Revised: Meanwhile, Λ( f ( x i ) , g k ( x i ) , y i , s ik ; Σ) is the conditional negative log probability of y i given s ik = 1 , i.e., − log p ( y i | s ik = 1 , x i ) .

Original: The specific form of Λ and miscellaneous model parameters Σ depends on the prediction task (either continuous-valued, binary, or multinomial outcome prediction).
Revised: The specific form of Λ and miscellaneous model parameters Σ depends on the prediction task (either continuous-valued, binary, or multinomial outcome prediction).

Original: For example, the objective function (8) for the continuous outcome is listed below in Equation 9, where we assume Y =f
Revised: For example, the objective function [Figure] (9) for the continuous outcome is listed below in Equation 10, where we assume Y =

Original:  ( X ) + ε ,
Revised: f ( X ) + ε ,

Original: S k
Revised: S k

Original: = I
Revised: = I

Original: [ g k ( X ) + η k > 0] , and [ η k , ε ] ∼ N (0 , [1 , σ ; ρ k ]) .
Revised: [ g k ( X ) + η k > 0] , and [ η k , ε ] ∼ N (0 , [1 , σ ; ρ k ]) .

Original: The details of its derivation and for other types of outcomes can be found in Appendix A.2.
Revised: The details of its derivation and for other types of outcomes can be found in Appendix A.2.

Original: [Equation]
Revised: [Equation]

Original: Ultimately, ˆ f is the outcome prediction model of primary interest.
Revised: Ultimately, ˆ f is the outcome prediction model of primary interest.

Original: For any input X from unseen domains, the HeckmanDG prediction is ˆ f ( X ) .
Revised: For any input X from unseen domains, the HeckmanDG prediction is ˆ f ( X ) .

Original: The proposed HeckmanDG differentiates from Heckman’s bias correction method in the following perspectives.
Revised: The proposed HeckmanDG differentiates from Heckman’s bias correction method in the following perspectives.

Original: First , it allows and models  domain-specific sample selection mechanisms.
Revised: First , it allows and models multiple domain-specific sample selection mechanisms.

Original: Second , by utilizing the multiple domains in the training data, HeckmanDG does not need auxiliary data[Figure] (features of examples from the target population but not sampled in the domain), which is a required input for the original Heckman model.
Revised: Second , by utilizing the multiple domains in the training data, HeckmanDG does not need auxiliary data (features of instances from the target population but not sampled in the domain), which is a required input for the original Heckman model.

Original: Third , HeckmanDG alleviates the linear assumptions in Heckman’s bias correction model to allow flexible forms for both prediction and selection functions including neural networks.
Revised: Third , HeckmanDG alleviates the linear assumptions in Heckman’s bias correction model to allow flexible forms for both prediction and selection functions including neural networks.

Original:  O PTIMIZATION
Revised:  O PTIMIZATION

Original: In this section, we consider the hypothesis sets F and G are neural networks, and present an efficient algorithm to optimize Equation 8.
Revised: In this section, we consider the hypothesis sets F and G are neural networks, and present an efficient algorithm to optimize Equation 9.

Original: Specifically, we propose a two-step approach that primarily trains g to optimum ( Step 1 ), then updates the remaining parameters of f and Σ ( Step 2 ).
Revised: Specifically, we propose a two-step approach that primarily trains g to optimum ( Step 1 ), then updates the remaining parameters of f and Σ ( Step 2 ).

Original: In Step 1, we learn the selection model parameters based on the following objective function:
Revised: In Step 1, we learn the selection model parameters based on the following objective function:

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: which essentially learns to predict the source domain memberships for training instances.
Revised: which essentially learns to predict the source domain memberships for training instances.

Original: In Step 2, we freeze the selection model ˆ g , and learn the remaining parameters of f and Σ .
Revised: In Step 2, we freeze the selection model ˆ g , and learn the remaining parameters of f and Σ .

Original: The proposed optimization algorithm is in part motivated by Heckman’s two-step estimator for bias correction, which was devised as a means of avoiding the nonlinearity of estimating both selection and outcome equations simulataneously (Heckman, 1979).
Revised: The proposed optimization algorithm is in part motivated by Heckman’s two-step estimator for bias correction, which was devised as a means of avoiding the nonlinearity of estimating both selection and outcome equations simulataneously (Heckman, 1979).

Original: We provide pseudocode describing the overall optimization procedure in Algorithm 1 in Appendix A.3, and an ablation study in Section 7 to support the necessity of two-step optimization.
Revised: We provide pseudocode describing the overall optimization procedure in Algorithm 1 in Appendix A.3, and an ablation study in Section 7 to support the necessity of two-step optimization.

Original: Neural Network Architecture.
Revised: Neural Network Architecture.

Original: We use a common feature extractor φ g : X → Z  for the domain selection models g
Revised: We use a common feature extractor φ g : X → Z g for the domain selection models g

Original: = { g k } Kk =1 , which only differ in the last linear predictors ω g,k : Z  → [0 , 1] , thus g k = ω g,k ◦ φ g .
Revised: = { g k } Kk =1 , which only differ in the last linear predictors ω g,k : Z g → [0 , 1] , thus g k = ω g,k ◦ φ g .

Revised: 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 φ g ( · ) .

Original: We define f = ω f ◦ φ f , where φ f : X → Z  is the feature extractor of the outcome prediction model and ω f : Z  → Y is the final linear predictor.
Revised: Similarly, we define f = ω f ◦ φ f , where φ f : X → Z f is the feature extractor of the outcome prediction model and ω f : Z f → Y is the final linear predictor.

Revised: In general, φ g and φ f are allowed to have different neural architectures with different numbers of parameters.

Revised: However, we found that simply using the same architecture (but learning different parameters) works well in practice.

Revised: Therefore, in all our experiments, HeckmanDG has roughly twice as much trainable parameters as other comparative methods.

Original: An overview of our neural network architecture is provided in Figure 1.
Revised: An overview of our neural network architecture is provided in Figure 2.

Original:   E XPERIMENTS6. S IMULATION
Revised: [Table]  E XPERIMENTS

Original: Model.
Revised: Simulation.

Original: We simulate a linear regression problem to assess HeckmanDG’s predictive performance.
Revised: We simulate a linear regression problem to assess HeckmanDG’s predictive performance.

Original: We simulate two covariates ( X 1 , X 2 ) and two training domains ( K = 2 ), and the following setting of domain selection and outcome mechanisms:
Revised: We simulate two covariates ( X 1 , X 2 ) and two training domains ( K = 2 ) based on the following setting of domain selection and outcome mechanisms:

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Equation]

Original: where we assume a normal prior over the selection coefficient α k 2 ∼ N ( µ α 2 , σ 2 α 2 ) in order to implicitly control the similarity between domains by differing its parameters.
Revised: where we assume a normal prior over the selection coefficient α k 2 ∼ N ( µ α 2 , σ 2 α 2 ) in order to implicitly control the similarity between domains by differing its parameters.

Original: We consider a true population of 100000 instances, from which we sample data for each domain with n k ∼ Uniform (1000 , 2000) .
Revised: We consider a true population of 100000 instances, from which we sample data for each domain with n k ∼ Uniform (1000 , 2000) .

Original: In each trial, we sample the selection coefficient α k 2 ∼ N (5 , 3 2 ) for the two training domains and a held-out in-distribution (ID) test set (similar to the training domains), and α k ′ 2 ∼ N ( − 5 , 3 2 ) for another held-out out-of-distribution (OOD) test set (dissimilar to the training domains).
Revised: In each trial, we sample the selection coefficient α k 2 ∼ N (5 , 3 2 ) for the two training domains and a held-out in-distribution (ID) test set (similar to the training domains), and α k ′ 2 ∼ N ( − 5 , 3 2 ) for another held-out out-of-distribution (OOD) test set (dissimilar to the training domains).

Original: We also simulate a random domain from the population.
Revised: We also simulate a random test set from the population.

Original: We assumed ρ k = 0 .
Revised: We assumed ρ k = 0 .

Original: 8 and σ = 1 .
Revised: 8 and σ = 1 .

Original:  Results.
Revised: Simulation Results.

Original: We observe that HeckmanDG not only outperforms ERM, but also other DG methods including IRM (Arjovsky et al., 2019), GroupDRO (Sagawa et al., 2019), and VREx (Krueger et al., 2021)(Table 1).
Revised: We observed that HeckmanDG not only outperforms ERM, but also other DG methods including IRM (Arjovsky et al., 2019), GroupDRO (Sagawa et al., 2019), and VREx (Krueger et al., 2021)(Table 1).

Original: Note that ‘ERM (oracle)’ is trained on iid training data, which serves as a theoretical lower bound on model performance.
Revised: Note that ‘ERM (oracle)’ is trained on iid training data, which serves as a theoretical lower bound on model performance.

Original: We highlight that not only does HeckmanDG performs well on the random test set, but also on the OOD test set.
Revised: We highlight that not only does HeckmanDG perform well on the random test set, but also on the OOD test set.

Original: In contrast, other methods tend to fit well to the train domains (ID performance is high), but generalize poorly to random and OOD test domains.
Revised: In contrast, other methods tend to fit well to the train domains (ID performance is high), but generalize poorly to random and OOD test domains.

Original: 6. B ENCHMARKS Datasets.
Revised: Benchmark Datasets.

Original: To further demonstrate the effectiveness of HeckmanDG on high-dimensional data regimes, we conducted experiments on four datasets from the W ILDS benchmark (Koh et al., 2021): 1) C AMELYON 17, 2) P OVERTY M AP , 3) I W ILD C AM , and 4) R X R X 1.
Revised: To further demonstrate the effectiveness of HeckmanDG on highdimensional data regimes, we conducted experiments on four datasets from the W ILDS benchmark (Koh et al., 2021): 1) C AMELYON 17, 2) P OVERTY M AP , 3) I W ILD C AM , and 4) R X R X 1.

Original: We used the same neural network architecture for selection and outcome feature extractors φ g and φ f , which are followed by linear predictors ω g and ω f .
Revised: We used the same neural network architecture for selection and outcome feature extractors φ g and φ f , which are followed by linear predictors ω g and ω f .

Original: For the domain selection model, we tuned hyperparameters to obtain the best domain selection model returning the highest macro F1 score on the training data, which in all cases achieve near-perfect accuracy.
Revised: For the domain selection model, we tuned hyperparameters to obtain the best domain selection model returning the highest macro F1 score on the training data, which in all cases achieve near-perfect accuracy.

Original: For the outcome prediction model, we adhered to the official guidelines and used the OOD validation set provided in the W ILDS repository for hyperparameter tuning and model selection based on the recommended metrics.
Revised: For the outcome prediction model, we adhered to the official guidelines and used the OOD validation set provided in the W ILDS repository for hyperparameter tuning and model selection based on the recommended metrics.

Original: Detailed descriptions of dataset statistics are presented in Table 5 of Appendix A.3.
Revised: Detailed descriptions of dataset statistics are presented in Table 6 of Appendix A.3.

Original: Details on hyperparameters and model training are presented in Table 6 of Appendix A.3.
Revised: Details on hyperparameters and model training are presented in Table 7 of Appendix A.3.

Original:  Results.
Revised: Benchmark Datasets Results.

Original: We summarized the results on the WILDS benchmark in Tables 2,3,4a, and 4b.
Revised: We summarized the results on the WILDS benchmark in Tables 2,3, and 4.

Original: All methods apply the same network architectures based on the W ILDS guideline for fair comparison.
Revised: All methods apply the same network architectures based on the W ILDS guideline for fair comparison.

Original: 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.
Revised: 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.

Original: We observed that HeckmanDG outperforms other methods on two out of four datasets, while performing on par with other methods on the remaining two.
Revised: We observed that HeckmanDG outperforms other methods on two out of four datasets, while performing on par with other methods on the remaining two.

Original: We make the following key observations.
Revised: We make the following key observations.

Original: First , HeckmanDG often robustly performs on the test domain although it may not generate the best performance on the validation domains.
Revised: First , HeckmanDG often robustly performs on the test domain although it may not generate the best performance on the validation domains.

Original: For example, on the C AMELYON 17 dataset, the performance gap between the validation and test datasets for HeckmanDG is 3 . 3 , which is substantially smaller than 14 . 
Revised: For example, on the C AMELYON 17 dataset, the performance gap between the validation and test datasets for HeckmanDG is 3 . 3 , which is substantially smaller than 14 . 1 of [Table]

Original: 1 of ‘ERM (scratch)’ and 7 . 1 of ‘ERM (ImageNet)’.
Revised: [Table] ‘ERM (scratch)’ and 7 . 1 of ‘ERM (ImageNet)’.

Original: This pattern is similar to what we observed in the simulation studies.
Revised: This pattern is similar to what we observed in the simulation studies.

Original: Because HeckmanDG is designed to predict for the underlying population, it may not produce the best prediction performances for the (non-randomly selected) source domains.
Revised: Because HeckmanDG is designed to predict for the underlying population, it may not produce the best prediction performances for the (non-randomly selected) source domains.

Original: 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.
Revised: 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.

Original: Second , HeckmanDG effectively works for both the domain shift and the subpopulation shift problems.
Revised: Second , HeckmanDG effectively works for both the domain shift and the subpopulation shift problems.

Original: This can be observed by looking at either the average or worst-group performance on the P OVERTY M AP dataset (i.e., 0.51). This [Table] [Table]
Revised: This can be observed by looking at either the average or worst-group performance on the P OVERTY M AP dataset (i

Original: [Table]
Revised: [Table]

Original:  again supports the robustness of HeckmanDG against a range of non-random selection probabilities of the testing domains.
Revised: e., 0.51). This again supports the robustness of HeckmanDG against a range of nonrandom selection probabilities of the testing domains.

Original:  A NALYSIS
Revised:  A NALYSIS

Original: Necessity of two-step optimization.
Revised: Necessity of two-step optimization.

Original: Recall that in Section 5, we have proposed a two-step optimization algorithm to learn the domain selection and outcome prediction models sequentially.

Original: To[Figure] (c) demonstrate the necessity and effectiveness of the proposed two-step optimization, we performed an ablation study on the P OVERTY M AP dataset, which has 13 training domains.
Revised: To demonstrate the necessity and effectiveness of the proposed two-step optimization, we performed an ablation study on the P OVERTY M AP dataset.

Original: For comparison, we trained g and f in one-step to simultaneously minimize Equation 8.
Revised: For comparison, we trained g and f in one-step to simultaneously minimize Equation 9.

Original: We kept the number of training epochs to 30 (with the same batch size) equivalent to make a valid comparison.

Original: We observed thatboth one-step and two-step optimization schemes were comparable  on training domains (Figure 2a), but yield different test performance on the OOD domains (Figure 2b).
Revised: We observed that one-step obtained comparable performances on training domains (Figure 3a), but yielded worse test performance on the OOD domains (Figure 3b).

Original: Compared to the one-step performances, the two-step optimization improved predictive performance on both average ( 0 . 78 → 0 . 81 ) and the worst-group ( 0 . 43 → 0 . 51 ).

Original: We believe thatpoor prediction of domain selection probabilities by a suboptimal ˆ g may mislead ˆ f to a deficient solution.
Revised: We believe that a suboptimal ˆ g may mislead ˆ f to a deficient solution in the one-step optimization (more details in Appendix).

Original: Relationship between the performance of g and f .
Revised: Relationship between the performance of g and f .

Original: To further investigate how the performance of a domain selection model influences the final outcome prediction, we trained f with selection models with different macro F1 scores, and compared their performance in Figure 2c.
Revised: To further investigate how the performance of a domain selection model influences the final outcome prediction, we take intermediate snapshots for ˆ g , and learned f based on each snapshot. On all five data split folds, we always observed a positive [Table] [Figure]

Original: Specifically, we take intermediate snapshots for ˆ g within a single training run to obtain selection models with varying performance, and learned f based on each snapshot.
Revised: [Table] [Figure]

Original: On all five data split folds, we always observed a positive correlation between the performances of ˆ g and ˆ f , demonstrating that a well-performing selection model is necessary to correctly guide the outcome prediction model to generalize to the underlying population.  C ONCLUSIONS
Revised: correlation between the performances of ˆ g and ˆ f (Figure 3c), demonstrating that a well-performing selection model is necessary to correctly guide the outcome prediction model.  C ONCLUSIONS

Original: We propose a Selection Guided Domain Generalization (SGDG) framework, in which we formulate domain generalization as a non-random sample selection problem and propose to jointly learn the prediction model f and the domain selection model g to achieve generalization on the true population.
Revised: We propose a Selection Guided Domain Generalization (SGDG) framework, in which we formulate domain generalization as a non-random sample selection problem and propose to jointly learn the prediction model f and the domain selection model g to achieve generalization on the true population.

Original: DG is a challenging problem as the particular structure of the distribution shift in the testing domains is unknown.
Revised: DG is a challenging problem as the particular structure of the distribution shift in the testing domains is unknown.

Original: In the presence of this uncertainty, we propose and theoretically justify the objective of minimizing the risk targeting the population distribution through SGDG.
Revised: In the presence of this uncertainty, we propose and theoretically justify the objective of minimizing the risk targeting the population distribution through SGDG.

Original: Furthermore, we have provided a set of Heckman-type SGDG estimators for various outcome types under parametric assumptions.
Revised: Furthermore, we have provided a set of Heckman-type SGDG estimators for various outcome types under parametric assumptions.

Original: Although it is still an open question if a single general-purpose 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.
Revised: Although it is still an open question if a single general-purpose 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.

Original: Note that SGDG can naturally utilize all domains in the training data, including the (outcome) labeled and unlabeled domains.
Revised: Note that SGDG can naturally utilize all domains in the training data, including the (outcome) labeled and unlabeled domains.

Original: The unlabeled domains will contribute to the estimation of ˆ g in
Revised: The unlabeled domains will contribute to the estimation of ˆ g in

Original: Equation 8, which indirectly improves generalization performance as we showed in Section 7.
Revised: Equation 9, which indirectly improves generalization performance as we showed in Section 7.

Original: An intriguing direction for future research is to explore whether we can improve SGDG performances by adapting ˆ g ( X ) for domains in the test data.
Revised: An intriguing direction for future research is to explore whether we can improve SGDG performances by adapting ˆ g ( X ) for domains in the test data.

Original: In that way, we may further refine the prediction ˆ f ( X ) for the unseen domains guided by their similarities to the source domains.
Revised: In that way, we may further refine the prediction ˆ f ( X ) for the unseen domains guided by their similarities to the source domains.

Original: A A PPENDIX A. P ROOFS
Revised: A A PPENDIX A. P ROOFS

Original: Theorem 1.
Revised: Theorem 1.

Original: The following chain of inequalities completes the proof:
Revised: The following chain of inequalities completes the proof:

Original: [Equation]
Revised: [Equation]

Original: A. D ETAILS OF H ECKMAN C ORRECTION -G UIDED D OMAIN G ENERALIZATION
Revised: A. D ETAILS OF H ECKMAN C ORRECTION -G UIDED D OMAIN G ENERALIZATION

Original: In this section
Revised: In this section

Original: [Equation]
Revised: [Equation]

Original: we present the exact form of our loss functions for continuous, binary [Equation] and multinomial outcomes, based on their specific parametric assumptions.
Revised: we present the exact form of our loss functions for continuous, binary, and multinomial outcomes, based on their specific parametric assumptions.

Original: Continuous Outcomes For regression tasks where y ∈ R , we assume that the selection and outcome is generated by the following data generation process:
Revised: Continuous Outcomes For regression tasks where y ∈ R , we assume that the selection and outcome is generated by the following data generation process:

Original: [Equation]
Revised: [Equation]

Original: where I ( · ) ∈ { 0 , 1 } is the indicator function, ρ k is the correlation between ρ k and ε , and σ ∈ R + is the standard deviation of ε .
Revised: where I ( · ) ∈ { 0 , 1 } is the indicator function, ρ k is the correlation between η k and ε , and σ ∈ R + is the standard deviation of ε .

Original: Being consistent with the notation in the main text, we hereby define Σ = { ρ 1 , . .
Revised: Being consistent with the notation in the main text, we hereby define Σ = { ρ 1 , . .

Original: . , ρ K , σ } .
Revised: . , ρ K , σ } .

Original: Denoting model parameters as θ = { f, g 1 , . .
Revised: Denoting model parameters as θ = { f, g 1 , . .

Original: . , g K , Σ } , we formulate the data likelihood as follows:
Revised: . , g K , Σ } , we formulate the data likelihood as follows:

Original: [Equation]
Revised: [Equation]

Original: and we take the negative log likelihood to formulate the loss function:
Revised: and we take the negative log likelihood to formulate the loss function:

Original: [Equation]
Revised: [Equation]

Original: Binary Outcomes For binary outcomes where y ∈ { 0 , 1 } , we assume a probit model for both selection and outcome:
Revised: Binary Outcomes For binary outcomes where y ∈ { 0 , 1 } , we assume a probit model for both selection and outcome:

Original: [Equation]
Revised: [Equation]

Original: where ρ k is the correlation between η k and ε .
Revised: where ρ k is the correlation between η k and ε .

Original: Herein, Σ
Revised: Herein, Σ

Original: = { ρ 1 , . .
Revised: = { ρ 1 , . .

Original: , ρ K } .
Revised: , ρ K } .

Original: Denoting model parameters as θ = { f, g 1 , . .
Revised: Denoting model parameters as θ = { f, g 1 , . .

Original: . , g K , Σ } , we formulate the data likelihood for binary outcomes as follows:
Revised: . , g K , Σ } , we formulate the data likelihood for binary outcomes as follows:

Original: [Equation]
Revised: [Equation]

Original: and take the negative log likelihood to define the loss function for binary outcomes:
Revised: and take the negative log likelihood to define the loss function for binary outcomes:

Original: [Equation]
Revised: [Equation]

Original: where Φ 2 ( · , · ; a ) is the cumulative density function of the bivariate standard normal given correlation a ∈
Revised: where Φ 2 ( · , · ; a ) is the cumulative density function of the bivariate standard normal given correlation a ∈

Original: [ − 1 , 1].
Revised: [ − 1 , 1]

Revised: [Equation]

Original: Multinomial Outcomes For multinomial outcome tasks where y ∈ { 1 , . . .
Revised: Multinomial Outcomes For multinomial outcome tasks where y ∈ { 1 , . . .

Original: , J } , we assume a multinomial probit outcome model (McFadden, 1989):
Revised: , J } , we assume a multinomial probit outcome model (McFadden, 1989):

Original: [Equation]
Revised: [Equation]

Original: 20) where [ η k , ε j ] ∼ N ( 0 , [1 , 1; ρ kj ]) and ρ kj is the correlation between ρ k and ε j .
Revised: [Equation] where [ η k , ε j ] ∼ N ( 0 , [1 , 1; ρ kj ]) and ρ kj is the correlation between ρ k and ε j .

Original: We further assume that the outcome error terms are independently distributed:
Revised: We further assume that the outcome error terms are independently distributed:

Original: ε = [ ε 1 , . .
Revised: ε = [ ε 1 , . .

Original: . , ε J ] ⊤ ∼
Revised: . , ε J ] ⊤ ∼

Original: N ( 0 , I ) .
Revised: N ( 0 , I ) .

Original: Consequently, Σ = { ρ 11 , . .
Revised: Consequently, Σ = { ρ 11 , . .

Original: , ρ KJ } .
Revised: , ρ KJ } .

Original: Denoting model parameters as θ = { f 1 , . . 
Revised: Denoting model parameters as θ = { f 1 , . . .

Original: , f J , g 1 , . . .
Revised: , f J , g 1 , . . .

Original: , g K , ρ 11 , . . .
Revised: , g K , ρ 11 , . . .

Original: , ρ KJ } , we formulate the multinomial outcome data likelihood as follows:
Revised: , ρ KJ } , we formulate the multinomial outcome data likelihood as follows:

Original: [Equation]
Revised: [Equation]

Original: where I [ · ] ∈ { 0 , 1 } is the indicator function.
Revised: where I [ · ] ∈ { 0 , 1 } is the indicator function.

Original: We use the negative logarithm of the data likelihood for the loss function:
Revised: We use the negative logarithm of the data likelihood for the loss function:

Original: [Equation]
Revised: [Equation]

Original: where T =
Revised: where T =

Original: [ − g k ( x i ) , ∞ (cid:1) × [ ξ j, 1 ( x i ) , ∞ (cid:1) × . . .
Revised: [ − g k ( x i ) , ∞ (cid:1) × [ ξ j, 1 ( x i ) , ∞ (cid:1) × . . .

Original: ×
Revised: ×

Original: [ ξ j,j − 1 ( x i ) , ∞ (cid:1) × [ ξ j,j +1 ( x i ) , ∞ (cid:1) × . . .
Revised: [ ξ j,j − 1 ( x i ) , ∞ (cid:1) × [ ξ j,j +1 ( x i ) , ∞ (cid:1) × . . .

Original: ×
Revised: ×

Original: [ ξ j,J ( x i ) , ∞ (cid:1) ⊆
Revised: [ ξ j,J ( x i ) , ∞ (cid:1) ⊆

Original: R J is the half-open J -dimensional hyperrectangular domain and ξ j,j ′ ( x i ) =
Revised: R J is the half-open J -dimensional hyperrectangular domain and ξ j,j ′ ( x i ) =

Original: − f j ( x i ) + f j ′ ( x i ) .
Revised: − f j ( x i ) + f j ′ ( x i ) .

Original: We use the GHK algorithm (Hajivassiliou & Ruud, 1994) to compute the multivariate normal integrals.
Revised: We use the GHK algorithm (Hajivassiliou & Ruud, 1994) to compute the multivariate normal integrals.

Original: A. O PTIMIZATION , D ATASETS , AND M ODEL T RAINING
Revised: A. O PTIMIZATION , D ATASETS , AND M ODEL T RAINING

Original: 7: end for 8: while g = { g k } Kk =1 Not Converged do 9: for all k = 1 , · · · , K do 10:

Original: B k ← BatchSampler ( D k , B ) 11: g k ← g k + γB ∇ (cid:80) ( x i ,s ik ) ∈B
Revised: B k ← BatchSampler ( D k , B ) 11: g k ← g k + γB ∇ (cid:80) ( x i ,s ik ) ∈B

Original: k [ s ik log Φ( g k ( x i ))
Revised: k [ s ik log Φ( g k ( x i ))

Original: + (1 − s ik ) log Φ( − g k ( x i ))] 12:end for 13: end while 14: ˆ g k = g k , for k = 1 , . . .
Revised: + (1 − s ik ) log Φ( − g k ( x i ))] 12:

Original: , K .
Revised: 14: ˆ

Original: 15: Step 2:Learn the Outcome Model ( f, Σ) 16:
Revised: 15: Step 2:

Original: D O
Revised: D O

Original: = D 17: while fNot Converged do 18:
Revised: D 17: while f Converged do 18:

Original: B= BatchSampler 
Revised: B BatchSampler (

Original: ( D O, B) 19: f ← f − γB ∇ (cid:80) ( x i ,y i , s i ) ∈B (cid:80)
Revised: O B 19: f ← f − γB ∇ (cid:80) ( x i ,y i , s i ) ∈B (cid:80)

Original: Kk =1 s ik Λ( f ( x i ) , ˆ g k ( x i ) , y i , s ik ; Σ) 20: Σ ← Σ − γB ∇ (cid:80) ( x i ,y i , s i ) ∈B (cid:80)
Revised: Kk =1 s ik Λ( f ( x i ) , ˆ g k ( x i ) , y i , s ik ; Σ) 20: Σ ← Σ − γB ∇ (cid:80) ( x i ,y i , s i ) ∈B (cid:80)

Original: Kk =1 s ik Λ( f ( x i ) , ˆ g k ( x i ), y i , s ik ; Σ) 21: end while 22:
Revised: Kk =1 s ik Λ( f ( x i ) , ˆ g k ( x i ) y i , s ik ; Σ 21: end while 22:

Original: ˆ f ← f 23: ˆ Σ ← Σ 24: returnˆ f
Revised: ˆ f ← f 23: ˆΣ ← Σ 24: return

Revised: [Table]

Original: [Table]
Revised: [Table]

Revised: training domain selection and outcome models.

Original: [Table]
Revised: Hyperparameters were determined through grid search on the OOD validation set. For optimizers we searched between { Adam , SGD } . Learning rates were searched among { 10 − 5 , 10 − 4 , 10 − 3 } , and weight decay among { 0 , 10 − 5 , 10 − 3 , 10 − 1 } . Note that since HeckmanDG training is two-phase, we searched for hyperparameters sequentially.

Revised: On all datasets, Σ is optimized with the Adam optimizer with a learning rate of 10 − 2 and no weight decay.

Revised: A. I MPLEMENTATION D ETAILS FOR SWAD AND L2A-OT

Revised: In Tables 2, 3, and 4, the numbers for SWAD (Cha et al., 2021) and L2A-OT

Revised: (Zhou et al., 2020) are reproduced based on our implementation.

Revised: For SWAD, we set the optimum patient parameter

Revised: N s = 3 , the overfitting patient parameter N e = 6 , and the tolerance rate r = 1 . 2 , which are the default values used in the original paper.

Revised: The evaluation frequency was set to 100.

Revised: On all four datasets, models were trained with SGD, using the same learning rates and weight decay factors as

Revised: ERM (Koh et al., 2021).

Revised: For L2A-OT, we set λ

Revised: Domain

Revised: = 0 .

Revised: 5 , λ Cycle = 10 , and λ CE = 1 .

Revised: On all three datasets (excluding P OVERTY M AP ), the generator G is trained with Adam using a constant learning rate of 3 · 10 − 4 , while the prediction model F is trained with SGD, using the same learning rates and weight decay factors as ERM (Koh et al., 2021).

Revised: We faithfully refer the readers to the original papers for the details regarding the hyperparameters.

Revised: [Figure]

Revised: A DDITIONAL A NALYSIS ON T WO -S TEP O PTIMIZATION

Original: [Table]
Revised: [Figure]