Selecting Treatment Effects Models for Domain Adaptation Using Causal Knowledge

1 INTRODUCTION . Causal inference models for estimating individualized treatment effects ( ITE ) are designed to provide actionable intelligence as part of decision support systems and , when deployed on mission-critical domains , such as healthcare , require safety and robustness above all ( Shalit et al. , 2017 ; Alaa & van der Schaar , 2017 ) . In healthcare , it is often the case that the observational data used to train an ITE model may come from a setting where the distribution of patient features is different from the one in the deployment ( target ) environment , for example , when transferring models across hospitals or countries . Because of this , it is imperative to select ITE models that are robust to these covariate shifts across disparate patient populations . In this paper , we address the problem of ITE model selection in the unsupervised domain adaptation ( UDA ) setting where we have access to the response to treatments for patients on a source domain , and we desire to select ITE models that can reliably estimate treatment effects on a target domain containing only unlabeled data , i.e. , patient features . UDA has been successfully studied in the predictive setting to transfer knowledge from existing labeled data in the source domain to unlabeled target data ( Ganin et al. , 2016 ; Tzeng et al. , 2017 ) . In this context , several model selection scores have been proposed to select predictive models that are most robust to the covariate shifts between domains ( Sugiyama et al. , 2007 ; You et al. , 2019 ) . These methods approximate the performance of a model on the target domain ( target risk ) by weighting the performance on the validation set ( source risk ) with known ( or estimated ) density ratios . However , ITE model selection for UDA differs significantly in comparison to selecting predictive models for UDA ( Stuart et al. , 2013 ) . Notably , we can only approximate the estimated counterfactual error ( Alaa & van der Schaar , 2019 ) , since we only observe the factual outcome for the received treatment and can not observe the counterfactual outcomes under other treatment options ( Spirtes et al. , 2000 ) . Consequently , existing methods for selecting predictive models for UDA that compute a weighted sum of the validation error as a proxy of the target risk ( You et al. , 2019 ) is suboptimal for selecting ITE models , as their validation error in itself is only an approximation of the model ’ s ability to estimate counterfactual outcomes on the source domain . To better approximate target risk , we propose to leverage the invariance of causal graphs across domains and select ITE models whose predictions of the treatment effects also satisfy known or discovered causal relationships . It is well-known that causality is a property of the physical world , and therefore the physical ( functional ) relationships between variables remain invariant across domains ( Schoelkopf et al. , 2012 ; Bareinboim & Pearl , 2016 ; RojasCarulla et al. , 2018 ; Magliacane et al. , 2018 ) . As shown in Figure 1 , we assume the existence of an underlying causal graph that describes the generating process of the observational data . We represent the selection bias present in the source observational datasets by arrows between the features { X1 , X2 } , and treatment T . In the target domain , we only have access to the patient features , and we want to estimate the patient outcome ( Y ) under different settings of the treatment ( intervention ) . When performing such interventions , the causal structure remains unchanged except for the arrows into the treatment node , which are removed . Contributions . To the best of our knowledge , we present the first UDA selection method specifically tailored for machine learning models that estimate ITE . Our ITE model selection score uniquely leverages the estimated patient outcomes under different treatment settings on the target domain by incorporating a measurement of how well these outcomes satisfy the causal relationships in the interventional causal graph GT . This measure , which we refer to as causal risk , is computed using a log-likelihood function quantifying the model predictions ’ fitness to the underlying causal graph . We provide a theoretical justification for using the causal risk , and we show that our proposed ITE model selection metric for UDA prefers models whose predictions satisfy the conditional independence relationships in GT and are thus more robust to changes in the distribution of the patient features . We also show experimentally that adding the causal risk to existing state-of-the-art model selection scores for UDA results in selecting ITE models with improved performance on the target domain . We provide an illustrative example of model selection for several real-world datasets for UDA , including ventilator assignment for COVID-19 . 2 RELATED WORKS . Our work is related to causal inference and domain adaptation . In this section , we describe existing methods for ITE estimation , UDA model selection in the predictive setting , and domain adaptation from a causal perspective . ITE models . Recently , a large number of machine learning methods for estimating heterogeneous ITE from observational data have been developed , leveraging ideas from representation learning ( Johansson et al. , 2016 ; Shalit et al. , 2017 ; Yao et al. , 2018 ) , adversarial training , ( Yoon et al. , 2018 ) , causal random forests ( Wager & Athey , 2018 ) and Gaussian processes ( Alaa & van der Schaar , 2017 ; 2018 ) . Nevertheless , no single model will achieve the best performance on all types of observational data ( Dorie et al. , 2019 ) and even for the same model , different hyperparameter settings or training iterations will yield different performance . ITE model selection . Evaluating ITE models ’ performance is challenging since counterfactual data is unavailable , and consequently , the true causal effects can not be computed . Several heuristics for estimating model performance have been used in practice ( Schuler et al. , 2018 ; Van der Laan & Robins , 2003 ) . Factual model selection only computes the error of the ITE model in estimating the factual patient outcomes . Alternatively , inverse propensity weighted ( IPTW ) selection uses the estimated propensity score to weigh each sample ’ s factual error and thus obtain an unbiased estimate ( Van der Laan & Robins , 2003 ) . Alaa & van der Schaar ( 2017 ) propose using influence functions to approximate ITE models ’ error in predicting both factual and counterfactual outcomes . Influence function ( IF ) based validation currently represents the state-of-the-art method in selecting ITE models . However , existing ITE selection methods are not designed to select models robust to distributional changes in the patient populations , i.e. , for domain adaptation . UDA model selection . UDA is a special case of domain adaptation , where we have access to unlabeled samples from the test or target domain . Several methods for selecting predictive models for UDA have been proposed ( Pan & Yang , 2010 ) . Here we focus on the ones that can be adapted for the ITE setting . The first unsupervised model selection method was proposed by Long et al . ( 2018 ) , who used Importance-Weighted Cross-Validation ( IWCV ) ( Sugiyama et al. , 2007 ) to select hyperparameters and models for covariate shift . IWCV requires that the importance weights ( or density ratio ) be provided or known ahead of time , which is not always feasible in practice . Later , Deep Embedded Validation ( DEV ) , proposed by You et al . ( 2019 ) , was built on IWCV by using a discriminative neural network to learn the target distribution density ratio to provide an unbiased estimation of the target risk with bounded variance . However , these proposed methods do not consider model predictions on the target domain and are agnostic of causal structure . Causal structure for domain adaptation . Recently , Kyono & van der Schaar ( 2019 ) proposed Causal Assurance ( CA ) as a domain adaptation selection method for predictive models that leverages prior knowledge in the form of a causal graph . Because their work is centered around predictive models , it is suboptimal for ITE models , where the edges into the treatment ( or intervention ) will capture the selection bias of the observational data . Furthermore , their method does not allow for examining the target domain predictions , which is a key novelty of this work . We leverage do-calculus ( Pearl , 2009 ) to manipulate the underlying directed acyclical graph ( DAG ) into an interventional DAG that more appropriately fits the ITE regime . More recently , researchers have focused on leveraging the causal structure for predictive models by identifying subsets of variables that serve as invariant conditionals ( Rojas-Carulla et al. , 2018 ; Magliacane et al. , 2018 ) . 3 PRELIMINARIES . 3.1 INDIVIDUALIZED TREATMENT EFFECTS AND MODEL SELECTION FOR UDA . Consider a training dataset Dsrc = { ( xsrci , tsrci , ysrci ) } Nsrci=1 consisting of Nsrc independent realizations , one for each individual i , of the random variables ( X , T , Y ) drawn from the source joint distribution pµ ( X , T , Y ) . Let pµ ( X ) be the marginal distribution of X . Assume that we also have access to a test dataset Dtgt = { xtgti } Ntgt i=1 from the target domain , consisting of Ntgt independent realizations of X drawn from the target distribution pπ ( X ) , where pµ ( X ) 6= pπ ( X ) . Let the random variable X ∈ X represent the context ( e.g . patient features ) and let T ∈ T describe the intervention ( treatment ) assigned to the patient . Without loss of generality , consider the case when the treatment is binary , such that T = { 0 , 1 } . However , note that our model selection method is also applicable for any number of treatments . We use the potential outcomes framework ( Rubin , 2005 ) to describe the result of performing an intervention t ∈ T as the potential outcome Y ( t ) ∈ Y . Let Y ( 1 ) represent the potential outcome under treatment and Y ( 0 ) the potential outcome under control . Note that for each individual , we can only observe one of potential outcomes Y ( 0 ) or Y ( 1 ) . We assume that the potential outcomes have a stationary distribution pµ ( Y ( t ) | X ) = pπ ( Y ( t ) | X ) given the context X ; this represents the covariate shift assumption in domain adaptation ( Shimodaira , 2000 ) . Observational data can be used to estimate E [ Y | X = x , T = t ] through regression . Assumption 1 describes the causal identification conditions ( Rosenbaum & Rubin , 1983 ) , such that the potential outcomes are the same as the conditional expectation : E [ Y ( t ) | X = x ] = E [ Y | X = x , T = t ] . Assumption 1 ( Consistency , Ignorability and Overlap ) . For any individual ( unit ) i , receiving treatment ti , we observe Yi = Y ( ti ) . Moreover , { Y ( 0 ) , Y ( 1 ) } and the data generating process p ( X , T , Y ) satisfy strong ignorability Y ( 0 ) , Y ( 1 ) ⊥⊥ T | X and overlap ∀x : P ( T | X = x ) > 0 . The ignorability assumption , also known as the no hidden confounders ( unconfoundedness ) assumptions , means that we observe all variables X that causally affect the assignment of the intervention and the outcome . Under unconfoundedness , X blocks all backdoor paths between Y and A ( Pearl , 2009 ) . Under Assumption 1 , the conditional expectation of the potential outcomes can also be written as the interventional distribution obtained by applying the do−operator under the causal framework of Pearl ( 2009 ) : E [ Y ( t ) | X = x ] = E [ Y | X = x , do ( T = t ) ] . This equivalence will enable us to reason about causal graphs and interventions on causal graphs in the context of selecting ITE methods for estimating potential outcomes . Evaluating ITE models . Methods for estimating ITE learn predictors f : X × T → Y such that f ( x , t ) approximates E [ Y | X = x , T = t ] = E [ Y ( t ) | X = x ] = E [ Y | X = x , do ( T = t ) ] . The goal is to estimate the ITE , also known as the conditional average treatment effect ( CATE ) : τ ( x ) = E [ Y ( 1 ) | X = x ] − E [ Y ( 0 ) | X = x ] ( 1 ) = E [ Y | X = x , do ( T = 1 ) ] − E [ Y | X = x , do ( T = 0 ) ] . ( 2 ) The CATE is essential for individualized decision making as it guides treatment assignment policies . A trained ITE predictor f ( x , t ) approximates CATE as : τ̂ ( x ) = f ( x , 1 ) − f ( x , 0 ) . Commonly used to assess ITE models is the precision of estimating heterogeneous effects ( PEHE ) ( Hill , 2011 ) : PEHE = Ex∼p ( x ) [ ( τ ( x ) − τ̂ ( x ) ) 2 ] , ( 3 ) which quantifies a model ’ s estimate of the heterogeneous treatment effects for patients in a population . UDA model selection . Given a set F = { f1 , . . . fm } of candidate ITE models trained on the source domain Dsrc , our aim is to select the model that achieves the lowest target risk , that is the lowest PEHE on the target domain Dtgt . Thus , ITE model selection for UDA involves finding : f̂ = arg min f∈F Ex∼pπ ( x ) [ ( τ ( x ) − τ̂ ( x ) ) 2 ] = arg min f∈F Ex∼pπ ( x ) [ ( τ ( x ) − ( f ( x , 1 ) − f ( x , 0 ) ) ) 2 ] . ( 4 ) For this purpose , we propose using the invariance of causal graphs across domains to select ITE predictors that are robust to distributional shifts in the marginal distribution of X .

Summary: the paper attacks the problem of model selection for individual treatement effect (ITE) models when the domain of learning and prediction differ. Proposal is to use causal consistency as an additional "regularizer" in existing domain adaptation (DA) model selection methods. The "regularizer" would be scoring to which extent replacing factual outcomes by their counterfactual predictions would preserve conditional independence relations (induced by the causal graph) in the prediction domain. Experiments on a variety of datasets are conducted to show the added performance induced by using the "regularizer".