To Impute or Not To Impute? Missing Data in Treatment Effect Estimation

1 INTRODUCTION . Treatment effects are arguably the most important estimand in causal inference [ Pearl , 2009 ; Neyman , 1923 ; Rubin , 1974 ] . A reason for this is that a treatment effect lies at the heart of a causal question . Using causal inference , we try to more explicitly attribute effect to the treatment in question , by carefully disentangling the role of the environment . We shall explain with a running example . Imagine , a job-training program to boost employment . Before sponsoring such a program , a legislative body may want to estimate its effect before widespread adoption . Using past data on the program , the body has to rely on causal methods to infer the effect of the training program before formal adoption . We define effect as the difference in outcome when applying the treatment versus not applying the treatment ( or any alternative treatment for that matter ) . Literature on inferring ( or predicting ) treatment effects is largely concerned with handling selection bias . That is , we identify a possible difference between the treated and non-treated subpopulations of the data , seeing that treatment is rarely distributed uniformly across the population . If not accounted for , selection bias will lead to biased estimates . As such , many works focus on novel strategies to handle this bias [ Johansson et al. , 2016 ; Alaa & van der Schaar , 2017 ; Shalit et al. , 2017 ; Rubin , 1974 ; Imbens & Rubin , 2015 ] . In the past , the program was not offered to everyone . Those with a job were less likely to be considered . As were those that had been unemployed for a long time . Perhaps age played a role ; being less likely to switch careers , older people would benefit less from job-training . Treatment effect models are adopted in a wide range of fields , such as : medicine [ Obermeyer & Emanuel , 2016 ; Alaa et al. , 2021 ; Bica et al. , 2021 ] , marketing [ Devriendt et al. , 2018 ; Ascarza , 2018 ; Debaere et al. , 2019 ] , or even human resources [ Rombaut & Guerry , 2020 ] . However , these methods have almost exclusively assumed that data is complete . From a practitioner ’ s standpoint , this may not always be the case and , in reality , data is often incomplete [ Burton & Altman , 2004 ; Lit- tle et al. , 2012 ] . When missingness is not properly accounted for when training these models , their adoption in critical environments , such as medicine , could lead to misguided decisions [ Garcı́aLaencina et al. , 2009 ; Little & Rubin , 2019 ] . We define two reasons why missingness may detrimentally influence treatment effects models . First , when certain covariates are missing when making a treatment decision , the fact that these variables are missing may also contribute to selection-bias . Without knowledge of someone ’ s age , it was deemed better to not offer the program . This is understandable , as the high price of the program requires some level of certainty . A second reason why missingness in the data may have disastrous effects on treatment effects models is that missingness may also be caused by the treatment . Missingness as a result of the treatment choice increases the difference between treatment subpopulations , even when there was no difference to begin with . Effectively , these set of covariates would be very much alike had it not been for their missingness . Adjusting on these missing variables introduces bias to the model . Once an applicant has accepted their offer for job-training , the organisers require some additional information . Imagine , the program ’ s registration process asking the applicant ’ s current address , the job of their spouse , the amount of children they have , etc . Had they not accepted to participate in the program , they would be less likely to provide this additional information . The literature proposes two solutions to handle missingness : ( i ) either we impute the missing variables before subjecting the data to a learner [ Rubin , 1978 ; 2004 ; Little & Rubin , 2019 ; Kallus et al. , 2018 ] ; or ( ii ) we consider a missing variable as another value , and use it directly [ Mayer et al. , 2020b ; a ; D ’ Agostino Jr & Rubin , 2000 ; Rosenbaum & Rubin , 1984 ] . We find that imputation effectively removes information— i.e . given that missingness may cause treatment-selection , it is much harder to handle selection bias when we are given only part of the total information . However , considering that missingness may be different between treatment subpopulations , it can also introduce bias where there is none to begin with . In this case , not imputing the data will introduce bias . Contribution . We introduce ( and motivate adoption of ) a formal description of missingness in data used to estimate treatment effects [ Neyman , 1923 ; Rubin , 1974 ] . In particular , we find that previous attempts— dating as far back as the 1980s [ Rosenbaum & Rubin , 1984 , Appendix B ] — at formalising missingness in treatment effects are too general and allow for inaccurate descriptions of missingness and its impact . We illustrate why these descriptions are insufficient , and provide an alternative termed mixed confounded missingness ( MCM ) . We argue that MCM is a general-purpose missingness mechanism distinct from well-known missingness mechansims such as missing ( completely ) at random [ Little & Rubin , 2019 ] , and a refinement of conditional independence of treatment [ Rosenbaum & Rubin , 1984 ] that should be adopted to describe missingness when estimating treatment effects . Furthermore , based on the insights provided by MCM , we propose a strategy to handling missing data in treatment effects , termed selective imputation . Our approach is theoretically motivated and we provide empirical evidence of how methods benefit from this approach and demonstrate the harm when missingness is not correctly dealt with . 2 PRELIMINARIES . Estimating causal effects is a difficult endeavor , as it requires us to answer a counterfactual question . In particular , when we observe the outcome after applying a treatment on an individual , it is impossible to also observe that individual ’ s outcome under alternative treatment [ Holland , 1986 ] . As a treatment effect is defined as the difference between both outcomes , we are tasked with inferring an estimand which is never observed , which is crucially different from standard supervised learning . One can estimate causal effects by conducting randomised controlled trials ( RCTs ) [ Fisher , 1925 ; Neyman , 1923 ] . However , RCTs are often very expensive , and are sometimes considered unethical in a clinical setting [ Hellman & Hellman , 1991 ; Edwards et al. , 1999 ] . However , the alternative we consider in our work— estimating effects from observational data —comes with its own challenges . Contrasting an RCT study , comparing the subpopulations associated with each treatment in an observational dataset will result in biased estimates . The reason lies in the difference between these subpopulations . If treatment is indeed not assigned randomly , but instead based on an individual ’ s characteristics , then these characteristics are more represented in each subpopulation as a result . If these characteristics also affect the outcome , then the outcomes become biased as the subpopulations are no longer comparable . This phenomenon is often termed selection bias . Notation . Let X ∈ X ⊆ Rd be the covariates of an individual ; let the individual be treated with W ∈ { 0 , 1 } ; and let Y ∈ Y ⊆ R be their observed outcome . Practically , X could be a patient with lung-cancer ; W = 1 could be chemo-therapy ( and W = 0 radio-therapy ) ; and Y their tumour size after treatment . We use a subscript , Xi to indicate the ith element in X , which means that Xi ∈ R. Assumptions in causal inference . Estimating unbiased treatment effects from observational data has received a lot of attention in recent years . One of the more popular avenues in the literature , is the potential outcomes ( POs ) framework of causality [ Neyman , 1923 ; Rubin , 1974 ] . We define the PO of a treatment w ∈ { 0 , 1 } as Y ( w ) , where Y ( w ) corresponds to the outcome an individual would have experienced had they been assigned treatment W = w. While the standard consistency assumption ( see Assum . 1 below ) allows us to interpret the observed outcome as the potential outcome of the observed treatment , i.e . Y = Y ( W ) ; selection bias makes estimating Y ( ¬W ) more involved . Countering selection bias is achieved by correctly adjusting for the confounders . In doing so , we make the following assumptions , standard in the PO-framework : Assumption 1 ( Consistency ) . The observed outcome Y = Y ( W ) = Y ( w ) if W = w , for w ∈ { 0 , 1 } and i = 1 , 2 , . . . , N , 1 i.e . outcomes in the data correspond to one of the potential outcomes . Assumption 2 ( Ignorability ) . The joint distribution p ( X , W , Y ) satisfies strong ignorability : Y ( 0 ) , Y ( 1 ) ⊥ W |X , i.e . the potential outcomes are independent of the treatment , conditioned on X , implying that there are no additional ( unobserved ) confounders beyond the variables in X . Assumption 3 ( Overlap ) . The distribution p ( X , W , Y ) satisfies overlap : ∃ δ ∈ ( 0 , 1 ) s.t . δ < p ( W |X = x ) < 1− δ , ∀x ∈ X , i.e . each individual has a probability to receive either treatment . Graphical models and causality . Alternatively to POs , we can express causal relationships as a graphical model . In particular , a causal relationship is depicted as a directed edge in a directed acyclic graph ( DAG ) [ Pearl , 2009 ] , where a parent node is the cause and the child node is the effect . The ignorability assumption in Assum . 2 is sometimes illustrated in such a graphical model [ Richardson & Robins , 2013 ] . Specifically , the ignorability assumption can be expressed as the DAG shown in Fig . 1a . Typically the influence of the treatment on the outcome is expressed as a single world intervention graph ( SWIG ) : W w Y [ Richardson & Robins , 2013 ] . We have removed this SWIG-path from our figures in order to focus our discussion on the path ( s ) between X and W . Note that the set of DAGs ( i.e . the Markov equivalence class ) that satisfy Assum . 2 encompasses more than just the DAG in Fig . 1a . Other DAGs in this equivalence class can just as easily respect Assum . 1 to 3 , but they would make no sense . For example , reversing the arrow between X and Y would imply that outcome causes the covariates in the individuals , while still respecting ignorability . Instead , Fig . 1a is motivated through logical reasoning , where treatment and outcome is caused by the covariates . Throughout the remainder of this paper , we will build ( and extend ) heavily on Fig . 1a . Causal estimands . We now arrive at our two estimands of interest : the average treatment effect ( ATE ) , and the conditional average treatment effect ( CATE ) . Given the notation above , we can define each estimand as follows : Definition 1 ( ATE ) . The ATE is defined as the population wide difference between a treatment ’ s potential outcomes . Mathematically , we can define the ATE as follows : τ̄ ( X ) : = EX [ Y ( 1 ) −Y ( 0 ) ] . Definition 2 ( CATE ) . The CATE is defined as a conditional difference between a treatment ’ s potential outcomes . Mathematically , we can define the CATE as follows : τ ( x ) : = E [ Y ( 1 ) −Y ( 0 ) |X = x ] . Missingness . In practice , a sample X may be incomplete . For example , a clinician responding to an urgent trauma case may have to select treatment based on incomplete information . In this scenario , the incomplete variables are considered missing . To learn from these data we could consider completing this sample through imputation , but missingness in itself may be informative . Perhaps the clinician ’ s decision would have been different if they had had complete information . Typically , we define three mechanisms to describe how a variable ended up to be missing : a first is missing completely at random ( MCAR ) , where missingness in one variable is independent of the other variables , a second is missing at random ( MAR ) , where the missingness in one variable may depend on 1The well-known stable unit treatment value assumption ( SUTVA ) assumes both no interference and consistency [ Rubin , 1980 ] . The equation in our consistency assumption also implies no interference . the other ( observed ) variables , and the third is missing not at random ( MNAR ) , where missingness is typically assumed to be caused by variables outside the observed covariates [ van Buuren , 2018 ; Rubin , 1976 ; Little & Rubin , 2019 ] . Generally , MCAR is attributed to noise in data-collection . We indicate the missing data in X with a variable Z ∈ { 1 , ? } d , where Zi = ? if Xi is missing , and Zi = 1 if Xi is observed . Having Z ⊥ X , Y ( w ) , corresponds with MCAR ; and Z 6⊥ X , but Z ⊥ Y ( w ) |X , corresponds with M ( N ) AR . In our paper , we denote the complete ( but unobserved ) sample as X , and the incomplete ( but observed ) sample as X̃ : = Z X ∈ { ? , R } d where we define as the element-wise product and we take Xi × ? = ? , i.e . if Xi is unobserved it equals ? in X̃ . Note that none of the existing missingness mechanisms take the treatment W into account as standard , since they span a broader literature beyond treatment effects . We have illustrated MCAR and M ( N ) AR as DAGs in Fig . 1b and Fig . 1c , respectively , where we have included arrows from X̃ and Z to W . This corresponds to the situation where treatment is decided based on what is actually observed ; the alternative where X is causing W would lead to a confounded setting , which is commonly assumed not to be the case in the literature . Note that these DAGs respect the conditional independent statements assumed in their respective missingness mechanisms . As we have noted in our introduction , we could either impute a missing valueXij with an estimate thereof , denoted Ẋij , such that we can predict a treatment effect from data with imputed samples ( i.e . learn τ̄ ( Ẋ ) or τ ( Ẋ ) ) ; or we could predict treatment effects from data with missingness directly ( i.e . learn τ̄ ( X̃ ) or τ ( X̃ ) ) .

Summary: In this paper, the authors mainly study the problem of missing data in treatment effect estimation and highlight the importance of addressing this problem. The authors propose a selective imputation scheme which is more well suited for addressing missingness in such scenarios. Authors also present several sample scenarios illustratively which indicate potential issues with current methods while doing TE with missingness.  Empirical results compare different scenarios where general imputation schemes (imputing all data, no data, wrong data) can be much worse than their proposed method.