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

Original: Anonymous authors

Original: 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 prove 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) es timators applicable to continuous binary and multinomial outcomes.

Original: We also demonstrate 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: In this paper, we consider DG as a non-random sample selection problem.
Revised: In this paper, we consider DG as a non-random sample selection problem.

Original: 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 .
Revised: 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: In a random sampling process, P ( S ki = 1) is independent and identically distributed ( iid ).
Revised: In a random sampling process, P ( S ki = 1) is independent and identically distributed ( iid ).

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

Original: Consequently, this leads to distributional shifts across domains: P jXY ̸ = P kXY for k ̸
Revised: Consequently, this leads to distributional shifts across domains: P jXY ̸ = P kXY for k ̸

Original: = j .
Revised: = j .

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

Original: 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 ).
Revised: We present a graphical model 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 ).

Original: In Figure 1, C 1 represents unobserved latent factors that correlate with X and S k , resulting in covariate sfhift.
Revised: In Figure 1, C 1 represents unobserved latent factors that correlate with X and S k , resulting in covariate shift.

Original: C 2 correlates with X , Y , and S simultaneously, entailing both covariate and concept shifts.
Revised: C 2 correlates with X , Y , and S simultaneously, entailing both covariate and concept shifts.

Original: 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. 
Revised: 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. [Equation] [Figure] [Equation]

Revised: [Figure]

Original: The vast majority of DG methods are developed to identify f that is robust to C 1 .
Revised: The vast majority of DG methods are developed to identify f that is robust to C 1 .

Original: However in practice, C 2 type of confounders often exist which make P ( S k = 1) related to both X and Y .For
Revised: However in practice, C 2 type of confounders often exist which make P ( S k = 1) related to both X and Y .

Original: [Figure] 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.
Revised: For 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.

Original: As a result, a model trained in an oncology specialist hospital may not be generalizable to a hospital serving veterans.
Revised: As a result, a model trained in an oncology specialist hospital may not be generalizable to a hospital serving veterans.

Original: Similarly, when we train a model to predict wealth index ( Y ) from satellite images ( X ) taken from different countries ( S k ), there may be latent factors such as the economic status ( C 2 ) correlating with X, Y and domains simultaneously.
Revised: Similarly, when we train a model to predict wealth index ( Y ) from satellite images ( X ) taken from different countries ( S k ), there may be latent factors such as the economic status ( C 2 ) correlating with X, Y and domains simultaneously.

Original: 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).
Revised: Therefore, a model trained on one country may not perform well in another country with a different rural/urban proportion or economic status (Koh et al., 2021).

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

Original: 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.
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) estimators 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.
Revised: DG has been studied under various contexts.

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

Original: Deng et al., 2020), entropy regularization (Zhao et al., 2020), and contrastive learning (Motiian et al., 2017; Kim et al., 2021).
Revised: Deng et al., 2020), entropy regularization (Zhao et al., 2020), and contrastive learning (Motiian et al., 2017; Kim et al., 2021).

Original: Other works exploit invariant causal effects across domains (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., 2021).

Original: Another family of robust optimization methods seek to minimize the worst-case error (Sagawa et al., 2019;
Revised: Another family of robust optimization methods seek to minimize the worst-case error (Sagawa et al., 2020;

Original: Xie et al., 2020; Krueger et al., 2021).
Revised: Xie et al., 2020; Krueger et al., 2021).

Original: More recently, other prominent directions of methods improve
Revised: More recently, other prominent directions of methods improve

Original: 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).
Revised: 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), 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 when 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 the availability of additional data drawn from 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: 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 .
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:  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 .

Original: Assumption 1 (Mutually Exclusive Domain Membership) .
Revised: Assumption 1 (Mutually Exclusive Domain Membership) .

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

Original: so that an instance can belong to one and only one domain.
Revised: so that an instance can belong to one and only one domain.

Original: Assumption 2 (Independent Domain Sampling Processes) .
Revised: Assumption 2 (Independent Domain Sampling Processes) .

Original: We assume that S k ⊥⊥ S j for all j ̸ = k .
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: A naive learning model 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 k XY 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 .

Original: Assumption 3 (Decomposable Loss Function) .
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: ℓ
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: f ( X ) , [Equation] 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, the negative log-likelihood under the probabilistic framework satisfies Assumption 3: 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: =
Revised: =

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 Assumption 3, 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 =1 ,min
Revised: { g k } Lk =1 ,

Original: [Equation]

Original: ∈F , g ∈G

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

Original: In such cases, we use P kX for k = 1 , · · · , K to learn the selection model g as in Problem (7).
Revised: In such cases, we use P kX for k = 1 , · · · , K to learn the selection model g as in Problem (7).

Original: Assumption 4.
Revised: Assumption 4.

Original: Let g ∗ be the optimal selection model that has the minimum expected ℓ s ,
Revised: Let g ∗ be the optimal selection model that has the minimum expected ℓ s ,

Original: [Equation]
Revised: [Equation]

Original: We assume that the first K coordinates ( g ∗ 1 , · · · , g ∗ K ) minimize (cid:80)Kk
Revised: We assume that the first K coordinates ( g ∗ 1 , · · · , g ∗ K ) minimize (cid:80)

Original: [Equation] 1 E X ∼P kX [ 
Revised: Kk =1 E X ∼P kX [ ℓ s ( g ( X ) , S k

Original: [Equation]
Revised: [Equation]

Original: s ( g ( X ) , S k = 1)]

Original: P ( S k = 1) .
Revised: P ( S k = 1) .

Original: 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: 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 } ).

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

Revised: [Equation]

Revised: , g ∈G

Original: [Equation]
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 7 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 ρ and the mean as a linear functions of X .
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: Building upon his work, we make the following assumption on the joint distribution of the outcome and the selection variables.
Revised: Building upon his work, we make the following assumption on the joint distribution of the outcome and the selection variables.

Original: Assumption 5 (Joint Distribution of Y (or latent variables ˜ Y ) and ˜ S ) .
Revised: Assumption 5 (Joint Distribution of Y (or latent variables ˜ Y ) and ˜ S ) .

Original: Let Y be the outcome variable and ˜ S
Revised: Let Y be the outcome variable and ˜ S

Original: k be the latent continuous variable where S k
Revised: k be the latent continuous variable where S k

Original: = I
Revised: = I

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

Original: We assume that
Revised: We assume that

Original: Y and S =
Revised: Y and S =

Original: ( 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 .
Revised: ( 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: For binary and multinomial outcomes, this assumption is on ˜ Y as the 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.

Original: Under this assumption, g = { g k } Kk =1 can be modeled as a set of independent probit models.
Revised: 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:
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
Revised: For example, the objective function

Original: [Figure]
Revised: [Figure]

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

Original: f ( 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 multiple 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 (features of instances 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 9.
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:

Original: [Equation] [Equation]
Revised: [Equation] [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 non-linearity of estimating both selection and outcome equations simultaneously (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 g 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 g → [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 .

Original: 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 ( · ) .
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: 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: 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.

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

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

Original: Therefore, in all our experiments, HeckmanDG has roughly twice as much trainable parameters as other comparative methods.
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 2.
Revised: An overview of our neural network architecture is provided in Figure 2.

Original: [Table]
Revised: [Table]

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: Simulation.
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 ) based on 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] [Equation]
Revised: [Equation] [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 test set 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: Simulation Results.
Revised: Simulation Results.

Original: We observed that HeckmanDG not only outperforms ERM, but also other
Revised: We observed that HeckmanDG not only outperforms ERM, but also other

Original: DG methods including IRM (Arjovsky et al., 2019), GroupDRO (Sagawa et al., 2019), and VREx (Krueger et al., 2021)(Table 1).
Revised: DG methods including IRM (Arjovsky et al., 2019), GroupDRO (Sagawa et al., 2020), 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 perform 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: Benchmark Datasets.
Revised: Benchmark Datasets.

Original: 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.
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 6 of Appendix A.3.
Revised: Detailed descriptions of dataset statistics are presented in Table 5 of Appendix A.3.

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

Original: Benchmark Datasets Results.
Revised: Benchmark Datasets Results.

Original: We summarized the results on the WILDS benchmark in Tables 2,3, and 4.
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
Revised: We observed that

Original: HeckmanDG outperforms other methods on two out of four datasets, while performing on par with other methods on the remaining two.
Revised: 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 .

Original: 1 of
Revised: 1 of

Original: [Table] [Table]
Revised: [Table] [Table]

Original: ‘ERM (scratch)’ and 7 . 1 of ‘ERM (ImageNet)’.
Revised: ‘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 ,
Revised: Second ,

Original: HeckmanDG effectively works for both the domain shift and the subpopulation shift problems.
Revised: 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
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: e., 0.51).
Revised: e., 0.51).

Original: This again supports the robustness of HeckmanDG against a range of nonrandom selection probabilities of the testing domains.
Revised: 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: To demonstrate the necessity and effectiveness of the proposed two-step optimization, we performed an ablation study on the P OVERTY M AP dataset.
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 8.

Original: We observed that one-step obtained comparable performances on training domains (Figure 3a), but yielded worse test performance on the OOD domains (Figure 3b).
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: We believe that a suboptimal ˆ g may mislead ˆ f to a deficient solution in the one-step optimization (more details in Appendix).
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 take intermediate snapshots for ˆ g , and learned f based on each snapshot.
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.

Original: On all five data split folds, we always observed a positive
Revised: On all five data split folds, we always observed a positive

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: correlation between the performances of ˆ g and ˆ f
Revised: correlation between the performances of ˆ g and ˆ f

Original: (Figure 3c), demonstrating that a well-performing selection model is necessary to correctly guide the outcome prediction model.
Revised: (Figure 3c), demonstrating that a well-performing selection model is necessary to correctly guide the outcome prediction model.

Original:  C ONCLUSIONS
Revised:  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 9, 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 for the unseen domains guided by their similarities to the source domains. A CKNOWLEDGEMENT

Revised: We thank the ICLR reviewers for their suggestions on improving the original version of this paper.

Revised: This work was supported in part by the National Institutes of Health under awards NIH R01LM013344, R01AG054467, R01AG065330, and R01-AG065330-02S1, by NYU Center for the

Revised: Study of Asian American Health under the NIH/NIMHD grant award #U54MD000538 (HD)

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 [Equation] we present the exact form of our loss functions for continuous, binary, and multinomial outcomes, based on their specific parametric assumptions.
Revised: In this section, 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 [Equation] ρ k is the correlation between η k and ε [Equation]

Revised: 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 [Equation] 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 , Σ } [Equation] 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] .

Original: [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):[Equation]
Revised: , J } , we assume a multinomial probit outcome model (McFadden, 1989):

Original: [Equation]
Revised: [Equation]

Original: where [ η k , ε j ] ∼ N ( 0 , [1 , 1; ρ kj ]) and ρ kj is the correlation between ρ k and ε j .
Revised: 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: 22) 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]

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

Revised: D

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

Original: k [ s ik log Φ( g k ( x i )) + (1 − s ik ) log Φ( − g k ( x i))] 12:
Revised: [ s ik log Φ( g k ( x i )) + (1 − s ik ) log Φ( − g k ( x i

Original: 14: ˆ

Original: 15: Step2:
Revised: 15: Step

Original: DO
Revised: D

Original:  D 17: whilef
Revised: = D 17: while

Original: Converged do 18:
Revised: Converged do 18:

Original: B
Revised: B

Original: BatchSampler(
Revised: BatchSampler

Original: O
Revised: D

Original: B
Revised: B

Original: 19: f ← f − γB ∇ (cid:80) ( x i ,y i , s i ) ∈B (cid:80)
Revised: 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 , sik ; Σ 21: end while 22: 
Revised: Kk =1 s ik Λ( f ( x i ) , ˆ g k ( x i ) , y i , s while 22: ˆ

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

Original: [Table] [Table]
Revised: [Table] [Table]

Original: training domain selection and outcome models.
Revised: domain selection and outcome models.

Original: Hyperparameters were determined through grid search on the OOD validation set.
Revised: Hyperparameters were determined through grid search on the OOD validation set.

Original: For optimizers we searched between { Adam , SGD } .
Revised: For optimizers we searched between { Adam , SGD } .

Original: Learning rates were searched among { 10 − 5 , 10 − 4 , 10 − 3 } , and weight decay among { 0 , 10 − 5 , 10 − 3 , 10 − 1 } .
Revised: Learning rates were searched among { 10 − 5 , 10 − 4 , 10 − 3 } , and weight decay among { 0 , 10 − 5 , 10 − 3 , 10 − 1 } .

Original: Note that since HeckmanDG training is two-phase, we searched for hyperparameters sequentially.
Revised: Note that since HeckmanDG training is two-phase, we searched for hyperparameters sequentially.

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

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

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

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

Original: For SWAD, we set the optimum patient parameter
Revised: For SWAD, we set the optimum patient parameter

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

Original: The evaluation frequency was set to 100.
Revised: The evaluation frequency was set to 100.

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

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

Original: For L2A-OT, we set λ
Revised: For L2A-OT, we set λ

Original: Domain
Revised: Domain

Original: = 0 .
Revised: = 0 .

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

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

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

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

Original: We provide further details on the analysis in Section 7, conducted to demonstrate the necessity and effectiveness of the proposed two-step optimization approach.
Revised: We provide further details on the analysis in Section 7, conducted to demonstrate the necessity and effectiveness of the proposed two-step optimization approach.

Original: 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, in order to make a valid comparison.
Revised: 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, in order to make a valid comparison.

Original: In Figure 4, we plot the training loss trajectories using different metrics.
Revised: In Figure 4, we plot the training loss trajectories using different metrics.

Original: We observed that g tends to converge to a suboptimal point if we optimize f , g , and Σ simultaneously by a one-step optimization approach (Figure 4a).
Revised: We observed that g tends to converge to a suboptimal point if we optimize f , g , and Σ simultaneously by a one-step optimization approach (Figure 4a).

Original: We suspect that this happens because the gradient of g (involving s ik = 1 ) depends on Σ , which changes during the joint onestep optimization process.
Revised: We suspect that this happens because the gradient of g (involving s ik = 1 ) depends on Σ , which changes during the joint onestep optimization process.

Original: This will often result in a suboptimal ˆg , as shown in Figure 4d.
Revised: This will often result in a suboptimal ˆg , as shown in Figure 4d.

Original: Note that the original Heckman correction (Heckman, 1979) also proposed a two-step approach to avoid the computational burden of having to estimate both f and g jointly.
Revised: Note that the original Heckman correction (Heckman, 1979) also proposed a two-step approach to avoid the computational burden of having to estimate both f and g jointly.

Original: One way to avoid this problem is to use an alternating optimization algorithm which alternately minimizes Equation 8 with respect to g , Σ , and f until convergence.
Revised: One way to avoid this problem is to use an alternating optimization algorithm which alternately minimizes Equation 8 with respect to g , Σ , and f until convergence.

Original: Our two-step approach can be regarded as a one-iteration alternating minimization procedure as it stops after a single iteration.
Revised: Our two-step approach can be regarded as a one-iteration alternating minimization procedure as it stops after a single iteration.

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

Original: However Equation 8 converged to a larger value by the one-step method, implying that it converged to a suboptimal point for g .
Revised: However Equation 8 converged to a larger value by the one-step method, implying that it converged to a suboptimal point for g .

Original: On the test datasets, the suboptimal solution of the onestep method resulted in worse domain generalization performance compared to the two-step method (Figure 3b).
Revised: On the test datasets, the suboptimal solution of the onestep method resulted in worse domain generalization performance compared to the two-step method (Figure 3b).

Original: [Figure]
Revised: [Figure]