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 paper proposes an algorithm that identifies disentangled representation to find out an individual treatment effect. A very specific model that tries to find out the underlying dynamics of such a problem is proposed and is learned by minimizing a suggested objective that takes the strengths of previous approaches. The method is demonstrated in a synthetic dataset and IHDP dataset and shown to outperform other previous methods by a large margin.