Targeted VAE: Structured Inference and Targeted Learning for Causal Parameter Estimation

Undertaking causal inference with observational data is extremely useful across a wide range of domains including the development of medical treatments , advertisements and marketing , and policy making . There are two main challenges associated with undertaking causal inference using observational data : treatment assignment heterogeneity ( i.e. , differences between the treated and untreated groups ) , and an absence of counterfactual data ( i.e . not knowing what would have happened if an individual who did get treatment , were instead to have not been treated ) . We address these two challenges by combining structured inference and targeted learning . To our knowledge , Targeted Variational AutoEncoder ( TVAE ) is the first method to incorporate targeted learning into deep latent variable models . Results demonstrate competitive and state of the art performance . 1 INTRODUCTION . The estimation of the causal effects of interventions or treatments on outcomes is of the upmost importance across a range of decision making processes and scientific endeavours , such as policy making ( Kreif & DiazOrdaz , 2019 ) , advertisement ( Bottou et al. , 2013 ) , the development of medical treatments ( Petersen et al. , 2017 ) , the evaluation of evidence within legal frameworks ( Pearl , 2009 ; Siegerink et al. , 2016 ) and social science ( Vowels , 2020 ; Hernan , 2018 ; Grosz et al. , 2020 ) . Despite the common preference for Randomized Controlled Trial ( RCT ) data over observational data , this preference is not always justified . Besides the lower cost and fewer ethical concerns , observational data may provide a number of statistical advantages including greater statistical power and increased generalizability ( Deaton & Cartwright , 2018 ) . However , there are two main challenges when dealing with observational data . Firstly , the group that receives treatment is usually not equivalent to the group that does not ( treatment assignment heterogeneity ) , resulting in selection bias and confounding due to associated covariates . For example , young people may prefer surgery , older people may prefer medication . Secondly , we are unable to directly estimate the causal effect of treatment , because only the factual outcome for a given treatment assignment is available . In other words , we do not have the counterfactual associated with the outcome for a different treatment assignment to that which was given . Treatment effect inference with observational data is concerned with finding ways to estimate the causal effect by considering the expected differences between factual and counterfactual outcomes . We seek to address the two challenges by proposing a method that incorporates targeted learning techniques into a disentangled variational latent model , trained according to the approximate maximum likelihood paradigm . Doing so enables us to estimate the expected treatment effects , as well as individual-level treatment effects . Estimating the latter is especially important for treatments that interact with patient attributes , whilst also being crucial for enabling individualized treatment assignment . Thus , we propose the Targeted Variational AutoEncoder ( TVAE ) , undertake an ablation study , and compare our method ’ s performance against current alternatives on two benchmark datasets . 2 BACKGROUND . Problem Formulation : A characterization of the problem of causal inference with no unobserved confounders is depicted in the Directed Acyclic Graphs ( DAGs ) shown in Figs . 1 ( a ) and 1 ( b ) . Fig . 1 ( a ) is characteristic of observational data , where the assignment of treatment is related to the covariates . Fig . 1 ( b ) is characteristic of the ideal RCT , where the treatment is unrelated to the covariates . Here , xi ∼ p ( x ) ∈ Rm represents the m-dimensional , pre-treatment covariates for individual i assigned factual treatment ti ∼ p ( t|x ) resulting in factual outcome yti ∼ p ( y|x , t ) . Together , these constitute dataset D = { [ yi , ti , xi ] } Ni=1 where N is the sample size . The conditional average treatment effect for an individual with covariates xi may be estimated as τ̂i ( xi ) = E [ yi|xi , do ( t = 1 ) − yi|xi , do ( t = 0 ) ] , where the expectation accounts for the nondeterminism of the outcome ( Jesson et al. , 2020 ) . Alternatively , by comparing the post-intervention distributions when we intervene on treatment t , the Average Treatment Effect ( ATE ) is τ̂ ( x ) = Ex [ E [ y|x , do ( t = 1 ) ] − E [ y|x , do ( t = 0 ) ] ] . Here , do ( t ) indicates the intervention on t , setting all instances to a static value , dynamic value , or distribution and therefore removing any dependencies it originally had ( Pearl , 2009 ; van der Laan & Rose , 2018 ; 2011 ) . This scenario corresponds with the DAG in Fig . 1 ( b ) , where treatment t is no longer a function of the covariates x . Using an estimator for the conditional mean Q ( t , x ) = E ( y|t , x ) , we can calculate the Average Treatment Effect ( ATE ) and the empirical error for estimation of the ATE ( eATE ) .1 In order to estimate eATE we assume access to the ground truth treatment effect τ , which is only possible with synthetic or semi-synthetic datasets . The Conditional Average Treatment Effect ( CATE ) may also be calculated and the Precision in Estimating Heterogeneous Effect ( PEHE ) is one way to evaluate a model ’ s efficacy in estimating this quantity . See the appendix for the complete definitions of these terms . The Naive Approach : The DAG in Fig . 1 ( a ) highlights the problem with taking a naive approach to modeling the joint distribution p ( y , t , x ) . The structural relationship t ← x → y indicates both that the assignment of treatment t is dependent on the covariates x , and that a backdoor path exists through x to y . In addition to our previous assumptions , if we also assume linearity , adjusting for this backdoor path is a simple matter of adjusting for x by including it in a logistic regression . The naive method is an example of the uppermost methods depicted in Fig . 2 , and leads to the largest bias . The problem with the approach is ( a ) that the graph is likely misspecified such that the true relationships between covariates as well as the relationships between covariates and the outcome may be more complex . There is also problem ( b ) , that linearity is not sufficient to ‘ let the data speak ’ ( van der Laan & Rose , 2011 ) or to avoid biased parameter estimates . Using powerful nonparametric models ( e.g. , neural networks ) may solve the limitations associated with linearity and interactions to yield a consistent estimator for p ( y|x ) , and such a model is an example of the middlemost methods depicted in Fig . 2 . However , this estimator is not targeted to the estimation of the causal effect parameter τ , only predicting the outcome , and we require a means to reduce residual bias . Targeted Learning : Targeted Maximum Likelihood Estimation ( TMLE ) ( Schuler & Rose , 2016 ; van der Laan & Rose , 2011 ; 2018 ; van der Laan & Starmans , 2014 ) falls under the lowermost 1For a binary outcome variable y ∈ { 0 , 1 } , E ( y|t , x ) is the same as the conditional probability distribution p ( y|t , x ) . methods depicted in Fig . 2 and follows an approach involving three main steps : ( 1 ) estimation of the conditional mean E ( y|t , x ) with estimatorQ0 ( t , x ) , ( 2 ) estimation of the propensity scores with estimator g ( t|x ) , and ( 3 ) updating the conditional mean estimatorQ0 to getQ∗ using the propensity scores to attain an estimate for the causal parameter τ . The propensity score for individual i is defined as the conditional probability of being assigned treatment g ( ti , xi ) = p ( t = ti|x = xi ) , ∈ [ 0 , 1 ] ( Rosenbaum & Rubin , 1983 ) . The scores can be used to compensate for the relationship between the covariates and the treatment assignment using Inverse Probability of Treatment Weights ( IPTWs ) , reweighting each sample according to its propensity score . Step ( 3 ) is undertaken using ‘ clever covariates ’ which are similar to the IPTWs . They form an additional covariate variable H ( 1 , xi ) = g ( 1|xi ) −1 for individual i assigned treatment , and H ( 0 , xi ) = −g ( 1|xi ) −1 for individual i not assigned treatment . Note that when we condition on a single numeric value we imply an intervention ( e.g . g ( 1|xi ) ≡ g ( do ( t = 1 ) |xi ) ) . A logistic regression is then undertaken as y = σ−1 [ Q0 ( t , x ) ] + H ( t , x ) where σ−1 is the logit/inverse sigmoid function , Q0 ( t , x ) is set to be a constant , suppressed offset and represents a fluctuation parameter which is to be estimated from the regression . Once has been estimated , we acquire an updated estimator : Q1 ( do ( t = t ) , x ) = σ [ σ−1 [ Q0 ( t , x ) ] + H ( t , x ) ] ( 1 ) This equation tells us that our new estimator Q1 is equal to the old estimator balanced by the corrective H ( t , x ) term . This term adjusts for the bias associated with the propensity scores . When the parameter is zero , it means that there is no longer any influence from the ‘ clever covariates ’ H ( ) . The updated estimator Q1 can then be plugged into the estimator for τ̂ ( Q1 ; x ) . When the optimal solution is reached ( i.e . when = 0 ) , the estimator also satisfies what is known as the efficient Influence Curve ( IC ) , or canonical gradient equation ( Hampel , 1974 ; van der Laan & Rose , 2011 ; Kennedy , 2016 ) : N∑ i=1 IC∗ ( yi , ti , xi ) = 0 = N∑ i=1 [ H ( ti , xi ) ( yi −Q ( ti , xi ) ) +Q ( 1 , xi ) −Q ( 0 , xi ) − τ ( Q ; x ) ] ( 2 ) where IC ( yi , ti , xi ) represents the IC , and IC∗ ( yi , ti , xi ) represents the efficient IC for consistent Q and g. It can be seen from the right hand side Eq . 2 that at convergence , the estimator and its estimand are equal : yi = Q ( ti , xi ) and Q ( 1 , xi ) − Q ( 0 , xi ) = τ ( Q ; x ) . Over the whole dataset , all terms in Eq . 2 ‘ cancel ’ resulting in the mean ¯IC = 0 . As such , the logistic regression in Eq . 1 represents a solution to the IC via a parametric submodel . The TMLE method provides a doubly robust , asymptotically efficient estimate of the causal or ‘ target ’ parameter , and these theoretical guarantees make it attractive for adaptation into neural networks for causal effect estimation .

Summary: The proposed contribution of this work is to build of the existing literature which uses variational autoencoders for causal inference by (1) allowing an explicit mechanism for modeling irrelevant covariates and (2) incorporating targeted regularization into the latent variable nnet framework. Optimization of the model is done by minimizing the ELBO subject to a penalty term which the authors refer to as “targeted regularization”. This term is essentially an application of the TMLE model. Both of these proposals appear to improve the performance on what are now pretty standard benchmark datasets (jobs and ihdp).