Learning Disentangled Representations for CounterFactual Regression

We consider the challenge of estimating treatment effects from observational data ; and point out that , in general , only some factors based on the observed covariates X contribute to selection of the treatment T , and only some to determining the outcomes Y . We model this by considering three underlying sources of { X , T , Y } and show that explicitly modeling these sources offers great insight to guide designing models that better handle selection bias in observational datasets . This paper is an attempt to conceptualize this line of thought and provide a path to explore it further . In this work , we propose an algorithm to ( 1 ) identify disentangled representations of the above-mentioned underlying factors from any given observational dataset D and ( 2 ) leverage this knowledge to reduce , as well as account for , the negative impact of selection bias on estimating the treatment effects from D. Our empirical results show that the proposed method achieves state-of-the-art performance in both individual and population based evaluation measures . 1 INTRODUCTION . As we rely more and more on artificial intelligence ( AI ) to automate the decision making processes , accurately estimating the causal effects of taking different actions gains an essential role . A prominent example is precision medicine – i.e. , the customization of health-care tailored to each individual patient – which attempts to identify which medical procedure t ∈ T will benefit a certain patient x the most , in terms of the treatment outcome y ∈ R. Learning such models requires answering counterfactual questions ( Rubin , 1974 ; Pearl , 2009 ) such as : “ Would this patient have lived longer [ and by how much ] , had she received an alternative treatment ? ” . For notation : a dataset D = { [ xi , ti , yi ] } Ni=1 used for treatment effect estimation has the following format : for the ith instance ( e.g. , patient ) , we have some context information xi ∈ X ⊆ RK ( e.g. , age , BMI , blood work , etc . ) , the administered treatment ti chosen from a set of treatment options T ( e.g. , { 0 : medication , 1 : surgery } ) , and the respective observed outcome yi ∈ Y ( e.g. , survival time ; Y ⊆ R+ ) as a result of receiving treatment ti . Note that D only contains the outcome of the administered treatment ( aka observed outcome : yi ) , but not the outcome ( s ) of the alternative treatment ( s ) ( aka counterfactual outcome ( s ) : yti for t ∈ T \ { ti } ) , which are inherently unobservable . For the binary-treatment case , we denote the alternative treatment as ¬ti = 1− ti . Pearl ( 2009 ) demonstrates that , in general , causal relationships can only be learned by experimentation ( on-line exploration ) , or running a Randomized Controlled Trial ( RCT ) , where the treatment assignment does not depend on the individual X – see Figure 1 ( a ) . In many cases , however , this is expensive , unethical , or even infeasible . Here , we are forced to approximate treatment effects from off-line datasets collected through Observational Studies . In such datasets , the administered treatment T depends on some or all attributes of individual X – see Figure 1 ( b ) . Here , as Pr ( T |X ) 6= Pr ( T ) , we say these datasets exhibit selection bias ( Imbens & Rubin , 2015 ) . Figure 2 illustrates selection bias in an example ( synthetic ) observational dataset . Here , we want to accurately estimate the Individual Treatment Effect ( ITE ) for each instance i – i.e. , to estimate ei = y1i − y0i . We frame the solution as learning the function f : X × T → Y that can accurately predict the outcomes ( both observed ŷiti as well as counterfactuals ŷi¬ti ) given the context information xi for each individual . As mentioned earlier , there are two challenges associated with estimating treatment effects : ( i ) The fact that counterfactual outcomes are unobservable ( i.e. , not present in any training data ) makes estimating treatment effects more difficult than the generalization problem in the supervised learning paradigm . This is an inherent characteristic of this task . ( ii ) Selection bias in observational datasets implies having fewer instances within each treatment arm at specific regions of the domain . This sparsity , in turn , would decrease the accuracy and confidence of predicting counterfactuals at those regions . This paper addresses the second challenge by investigating the root causes of selection bias , by dissecting and identifying the underlying factors that can generate an observational dataset D , and leveraging this knowledge to reduce , as well as account for , the negative impact of selection bias on estimating the treatment effects from D. In this work , we borrow ideas from the representation learning literature ( Bengio et al. , 2013 ) in order to reduce selection bias and from the domain adaptation literature ( Shimodaira , 2000 ) in order to account for the remainder selection bias that ( might ) still exist after its reduction . Our analysis relies on the following assumptions : Assumption 1 : Unconfoundedness ( Rosenbaum & Rubin , 1983 ) – There are no unobserved confounders ( i.e. , covariates that contribute to both treatment selection procedure as well as determination of outcomes ) . Formally , { Y t } t∈T ⊥ T | X . Assumption 2 : Overlap ( Imbens , 2004 ) – Every individual x should have a non-zero chance of being assigned to any treatment arm . That is , Pr ( T = t |X=x ) 6= 0 ∀t ∈ T , ∀x ∈ X . These two assumptions together are called strong ignorability ( Rosenbaum & Rubin , 1983 ) . Imbens & Wooldridge ( 2009 ) showed that strong ignorability is sufficient for ITE to be identifiable . Without loss of generality , we assume that the random variable X follows a ( n unknown ) joint probability distribution Pr ( X |Γ , ∆ , Υ ) , treatment T follows Pr ( T |Γ , ∆ ) , and outcome Y T follows Pr T ( Y T |∆ , Υ ) , where Γ , ∆ , and Υ represent the three underlying factors1 that generate an obser1 Examples for : ( Γ ) rich patients receiving the expensive treatment while poor patients receiving the cheap one – although outcomes of the possible treatments are not particularly dependent on patients ’ wealth status ; ( ∆ ) younger patients receiving surgery while older patients receiving medication ; and ( Υ ) genetic information that determines the efficacy of a medication , however , such relationship is unknown to the attending physician . vational dataset D. The respective graphical model is illustrated in Figure 3 . Conforming with the statements above , note that the graphical model also suggests that selection bias is induced by factors Γ and ∆ , where ∆ represents the confounding factors between T and Y . Main contribution : We argue that explicit identification of the underlying factors { Γ , ∆ , Υ } in observational datasets offers great insight to guide designing models that better handle selection bias and consequently achieve better performance in terms of estimating ITEs . In this paper , we propose a model , named Disentangled Representations for CounterFactual Regression ( DR-CFR ) , that is optimized to do exactly that . We also present experiments that demonstrate the advantages of this perspective ; and show empirically that the proposed method outperforms state-of-the-art models in a variety of data generation scenarios with different dimensionality of factors ; see below . 2 RELATED WORKS . Selection bias in observational datasets is equivalent to a domain adaptation scenario where a model is trained on a “ source ” ( observed ) data distribution , but should perform well on a “ target ” ( counterfactual ) one . Learning treatment effects from observational datasets is closely related to “ off-policy learning from logged bandit feedback ” – cf. , ( Swaminathan & Joachims , 2015a ) , whose goal is learning an optimal policy that selects the best personalized treatment for each individual . A common statistical solution is re-weighting certain data instances to balance the source and target distributions . The majority of re-weighting approaches belong to the Inverse Propensity Weighting ( IPW ) family of methods – cf. , ( Austin , 2011 ; Bottou et al. , 2013 ; Swaminathan & Joachims , 2015c ) . While IPW methods are unbiased , they suffer from high variance . Swaminathan & Joachims ( 2015b ) proposed the Counterfactual Risk Minimization ( CRM ) principle to alleviate this issue . In summary , re-weighting is an attempt to account for the selection bias . Johansson et al . ( 2016 ) is among the pioneer works that explored ways to use techniques from representation learning ( Bengio et al. , 2013 ) to reduce the selection bias . Shalit et al . ( 2017 ) present a refined version of ( Johansson et al. , 2016 ) ’ s method that learns a common representation space Φ ( x ) = φ by minimizing the discrepancy ( Mansour et al. , 2009 ) ( hereinafter “ disc ” ) between the conditional distributions of φ given t=0 versus φ , given t=1 . That is , disc ( { Φ ( xi ) } i : ti=0 , { Φ ( xi ) } i : ti=1 ) ( 1 ) which is ( effectively ) a regularization term that attempts to reduce selection bias in the learned representation . On top of this representation learning network , they trained two regression networks ht ( φ ) – one for each treatment arm ( t ∈ { 0 , 1 } ) – that predict the outcomes . Hassanpour & Greiner ( 2019 ) argued that the learned representation can not and should not remove all the selection bias , as the confounders not only contribute to choosing a treatment but also to determining the respective outcomes.2 As a result , where there are confounders ( which is a common situation ) , even φ would exhibit some selection bias , although less than that in the original domain x . They built on the work of ( Shalit et al. , 2017 ) by introducing context-aware importance sampling 2 While Hassanpour & Greiner ( 2019 ) presented a graphical model similar to our Figure 3 , they only used it to investigate the nature of selection bias . N.b. , they did not implement the idea of learning disentangled representations for counterfactual regression ; instead , their method [ like ( Shalit et al. , 2017 ) ] learns a common representation φ that can represent only the confounders , but not the other factors . Our approach extends theirs by providing an algorithm that can learn disentangled representations of the underlying factors from observational datasets . weights , that attempt to account for the above-mentioned remainder selection bias . These weights ωi = 1 + Pr ( φi | ¬ti ) Pr ( φi | ti ) = 1 + Pr ( ti ) 1− Pr ( ti ) · 1− π ( ti |φi ) π ( ti |φi ) ( 2 ) are designed to enhance performance of estimating both factual as well as counterfacual outcomes ( by the 1 and Pr ( φ | ¬t ) Pr ( φ | t ) terms , respectively ) , where π ( ti |φi ) is the probability of assigning the observed ti conditioned on the learned context φi . Note that both ( Shalit et al. , 2017 ) and ( Hassanpour & Greiner , 2019 ) use Φ to model the concatenation of factors ∆ and Υ ( see Figure 3 ) . Although it does make sense that there should be no discrepancy between conditional distributions of Υ , the ∆ factor should model the confounding factors , which by definition , must embed some information about treatment assignment . This would result in a positive discrepancy between conditional distributions of ∆ that should not be minimized . Thus , minimizing Equation ( 1 ) with respect to Φ can lead to problematic results as it discards some of the confounders . Yao et al . ( 2018 ) proposed the Similarity preserved Individual Treatment Effect ( SITE ) method , which extends Shalit et al . ( 2017 ) ’ s framework by adding a local similarity preserving component . This component acts as a regularization term , that attempts to retain the same neighbourhood relationships in the learned representation space as exhibited in the original space , by matching the propensity scores Pr ( t=1 |x ) and Pr ( t=1 |φ ) . This , however , results in learning sub-optimal representations when Γ 6= ∅ as SITE tries to keep instances whose Γs are far apart , also far apart in φ . In other words , this component penalizes reducing selection bias in φ by not discarding the irrelevant information present in Γ even when it does not hurt the outcome estimation at all . Our work has many similarities to ( Kuang et al. , 2017 ) , who decomposed X into two subsets : confounding and adjustment variables , which are similar to our ∆ and Υ factors respectively . They then used an optimization algorithm for identifying these variables , to ultimately find an unbiased estimate of the Average Treatment Effect ( ATE ) . We extend their work in three ways : ( i ) In addition to confounders and adjustment variables , we also identify the factors that determine the treatment and have no effect on the outcome ( i.e. , Γ ) . ( ii ) Unlike ( Kuang et al. , 2017 ) that take a linear approach by tagging the raw features as either confounders or adjustment variables , our proposed method has the capacity to learn [ non-linear ] representations of the underlying factors . ( iii ) Our method facilitates estimating both ATE as well ITE , whereas ( Kuang et al. , 2017 ) can not provide estimates of ITEs .

Summary:    The authors consider the problem of estimating average treatment effects when observed X and treatment T causes Y. Observational data for X,T,Y is available and strong ignorability is assumed. Previous work (Shalit et al 2017) introduced learning a representation that is invariant in distribution across treatment and control groups and using that with treatment to estimate Y. However, authors point out that this representation being forced to be invariant still does not drive the selection bias to zero. A follow up work (Hassanpour and Greiner 2019) - corrects for this by using additional importance weighting that estimates the treatment selection bias given the learnt representation. However, the authors point out even this is not complete in general, as X could be determined by three latent factors, one that is the actual confounder between treatment and outcome and the other that affects only the outcome and the other that affects only the treatment. Therefore, the authors propose to have three representations and enforce independence between representation that solely determines outcome and the treatment and make other appropriate terms depend on the respective latent factors. This gives a modified objective with respect to these two prior works.