Variational Auto-Encoder Architectures that Excel at Causal Inference

1 INTRODUCTION . As one of the main tasks in studying causality ( Peters et al. , 2017 ; Guo et al. , 2018 ) , the goal of Causal Inference is to figure out how much the value of a certain variable would change ( i.e. , the effect ) had another certain variable ( i.e. , the cause ) changed its value . A prominent example is the counterfactual question ( Rubin , 1974 ; Pearl , 2009 ) “ Would this patient have lived longer [ and by how much ] , had she received an alternative treatment ? ” . Such question is often asked in the context of precision medicine , 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 ( e.g. , survival time ) . A fundamental problem in causal inference is the unobservablity of the counterfactual outcomes ( Holland , 1986 ) . That is , for each subject i , any real-world dataset can only contain the outcome of the administered treatment ( aka the observed outcome : yi ) , but not the outcome ( s ) of the alternative treatment ( s ) ( aka the counterfactual outcome ( s ) ) — i.e. , yti for t ∈ T \ { ti } . In other words , the causal effect is never observed ( i.e. , missing in any training data ) and can not be used to train predictive models , nor can it be used to evaluated a proposed model . This makes estimating causal effects a more difficult problem than that of generalization in the supervised learning paradigm . In general , we can categorize most machine learning algorithms into two general approaches , which differ in how the input features x and their target values y are modeled ( Ng & Jordan , 2002 ) : Discriminative methods focus solely on modeling the conditional distribution p ( y|x ) with the goal of direct prediction of y for each instance x . For prediction tasks , discriminative approaches are often more accurate since they use the model parameters more efficiently than generative approaches . Most of the current causal inference methods are discriminative , including the Balancing Neural Network ( BNN ) ( Johansson et al. , 2016 ) , CounterFactual Regression Network ( CFR-Net ) ( Shalit et al. , 2017 ) , and CFR-Net ’ s extensions — cf. , ( Yao et al. , 2018 ; Hassanpour & Greiner , 2019 ; 2020 ) — as well as Dragon-Net ( Shi et al. , 2019 ) . Generative methods , on the other hand , describe the relationship between x and y by their joint probability distribution p ( x , y ) . This , in turn , would allow the generative model to answer arbitrary queries , including coping with missing features x using the marginal distribution p ( x ) or [ similar to discriminative models ] predicting the unknown target values y via p ( y|x ) . A promising direction forward for causal inference is developing generative models , using either Generative Adverserial Network ( GAN ) ( Goodfellow et al. , 2014 ) or Variational Auto-Encoder ( VAE ) ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) . This has led to two generative approaches for causal inference : GANs for inference of Individualised Treatment Effects ( GANITE ) ( Yoon et al. , 2018 ) and Causal Effect VAE ( CEVAE ) Louizos et al . ( 2017 ) . However , neither of the two achieve competitive performance in terms of treatment effect estimation compared to the discriminative approaches . Although discriminative models have excellent predictive performance , they suffer from two drawbacks : ( i ) overfitting , and ( ii ) making highly-confident predictions , even for instances that are “ far ” from the observed training data . Generative models based on Bayesian inference , on the other hand , can handle both of these drawbacks : issue ( i ) can be minimized by taking an average over the posterior distribution of model parameters ; and issue ( ii ) can be addressed by explicitly providing model uncertainty via the posterior ( Gordon & Hernández-Lobato , 2020 ) . Although the exact inference is often intractable , efficient approximations to the parameter posterior distribution is possible through variational methods . Here , we use the Variational Auto-Encoder ( VAE ) ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) for the Bayesian inference component of our causal inference method . Contribution : In this paper , we propose three interrelated Bayesian model architectures ( namely Series , Parallel , and Hybrid ) that employ the VAE framework to address the task of causal inference for binary treatments . We find that the best performing architecture is the Hybrid model , that is [ partially ] successful in decomposing the underlying factors of any observational dataset . This is a valuable property , as that means it can accurately estimate all all treatment outcomes . We demonstrate that these models significantly outperform the state-of-the-art in terms of treatment effect estimation performance on two publicly available benchmarks , as well as a fully synthetic dataset that allows for detailed performance analyses . 2 RELATED WORKS . CFR-Net Shalit et al . ( 2017 ) considered the binary treatment task and attempted to learn a representation space Φ that reduces selection bias by making Pr ( Φ ( x ) | t=0 ) and Pr ( Φ ( x ) | t=1 ) as close to each other as possible , provided that Φ ( x ) retains enough information that the learned regressors { ht ( Φ ( · ) ) : t∈ { 0 , 1 } } can generalize well on the observed outcomes . Their objective function includes L [ yi , h ti ( Φ ( xi ) ) ] , which is the loss of predicting the observed outcome for sample i ( described as xi ) , weighted by ωi = ti2u + 1−ti 2 ( 1−u ) , where u = Pr ( t=1 ) . This is effectively setting ωi = 1 2 Pr ( ti ) where Pr ( ti ) is the probability of selecting treatment ti over the entire population . DR-CFR Hassanpour & Greiner ( 2020 ) argued against the standard implicit assumption that all of the covariates X are confounders ( i.e. , contribute to both treatment assignment and outcome determination ) . Instead , they proposed a graphical model similar to that in Figure 1 and designed a discriminative causal inference approach accordingly — built on top of the CFR-Net . Specifically , their model , named Disentangled Representations for CFR ( DR-CFR ) , includes three representation networks , each trained with constraints to insure that each component corresponds to its respective underlying factor . While the idea behind DR-CFR provides an interesting intuition , it is known that only generative models ( and not discriminative ones ) can truly identify the underlying data generating mechanism . This paper is a step in this direction . Dragon-Net Shi et al . ( 2019 ) ’ s main objective was to estimate the Average Treatment Effect ( ATE ) , which they explain requires a two stage procedure : ( i ) fit models that predict the outcomes for both treatments ; and ( ii ) find a downstream estimator of the effect . Their method is based on a classic result from strong ignorability — i.e. , Theorem 3 in ( Rosenbaum & Rubin , 1983 ) — that states : ( y1 , y0 ) ⊥ t |x & Pr ( t = 1 |x ) ∈ ( 0 , 1 ) =⇒ ( y1 , y0 ) ⊥ t | b ( x ) & Pr ( t = 1 | b ( x ) ) ∈ ( 0 , 1 ) where b ( x ) is a balancing score1 . They consider propensity score as a balancing score and argue that only the parts ofX relevant for predicting T are required for the estimation of the causal effect 2 . This theorem only provides a way to match treated and control instances though — i.e. , it helps finding potential counterfactuals from the alternative group to calculate ATE . Shi et al . ( 2019 ) , however , used this theorem to derive minimal representations on which to regress to estimate the outcomes . 1That is , X ⊥ T | b ( X ) ( Rosenbaum & Rubin , 1983 ) . 2The authors acknowledge that this would hurt the predictive performance for individual outcomes . As a result , this yields inaccurate estimation of Individual Treatment Effects ( ITEs ) . GANITE Yoon et al . ( 2018 ) proposed the counterfactual GAN , whose generator G , given { x , t , yt } , estimates the counterfactual outcomes ( ŷ¬t ) ; and whose discriminator D tries to identify which of { [ x , 0 , y0 ] , [ x , 1 , y1 ] } is the factual outcome . It is , however , unclear why this requires that G must produce samples that are indistinguishable from the factual outcomes , especially as D can just learn the treatment selection mechanism instead of distinguishing the factual outcomes from counterfactuals . Although this work is among the few generative approaches for causal inference , our empirical results ( in Section 4 ) show that it does not effectively estimate counterfactual outcomes . CEVAE Louizos et al . ( 2017 ) used VAE to extract latent confounders from their observed proxies in X . While this is an interesting step in the right direction , empirical results show that it does not always accurately estimate treatment effect ( see Section 4 ) . The authors note that this may be because CEVAE is not able to address the problem of selection bias . Another reason that we think contributes to CEVAE ’ s sub-optimal performance is its assumed graphical model of the underlying data generating mechanism ( depicted in Figure 2 ) . This model assumes that there is only one latent variable Z ( confounding T and Y ) that generates the entire observational data ; however , we know from ( Kuang et al. , 2017 ) and ( Hassanpour & Greiner , 2020 ) that there must be more ( see Figure 1 ) . R2 : M1 and M2 VAEs In an attempt to enhance the conventional representation learning with VAEs — referred to as the M1 model ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) — in a semisupervised manner , Kingma et al . ( 2014 ) proposed the M2 VAE . While the M1 model helps learning latent representations from the covariate matrix X alone , the M2 model allows the target information also to guide the representation learning process . In our work , the target information includes the treatment bit T as well as the observed outcome Y . This additional information helps learning more expressive representations , that was not possible with the unsupervised M1 model . Appendix A.1 presents a more detailed overview of the M1 and M2 VAEs . 3 METHOD . Following ( Hassanpour & Greiner , 2020 ) and without loss of generality , we assume that the random variable X follows an unknown joint probability distribution Pr ( X |Γ , ∆ , Υ , Ξ ) , where Γ , ∆ , Υ , and Ξ are non-overlapping independent factors . Moreover , we assume that treatment T follows Pr ( T |Γ , ∆ ) ( i.e. , Γ and ∆ are the responsible factors for selection bias ) and outcome Y T follows Pr T ( Y T |∆ , Υ ) ; see Figure 1 . Observe that the factor Γ ( resp. , Υ ) partially determines only T ( resp. , Y ) , but not Y ( resp. , T ) ; and ∆ includes the confounding factors between T and Y . Our goal is to design generative model architectures that encourage learning disentangled representations of these four underlying latent factors ( see Figure 1 ) . In other words , it is an attempt to decompose and separately learn the underlying factors that are responsible for determining T and Y . To achieve this , we propose three architectures ( as illustrated in Figures 3 ( a ) , 3 ( b ) , and 3 ( c ) ) , each employing a VAE ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) that each include a decoder ( generative model ) and an encoder ( variational posterior ) . Specifically , we use the M1 and M2 models from ( Kingma et al. , 2014 ) as our building blocks , leading to a Series architecture , a Parallel architecture , and a Hybrid one . Each component is parametrized as a deep neural network .

Summary: Some generative models have been proposed for causal effect estimation but they often do not have a competitive performance. Recent work suggested that a combination of generative and discriminative model may improve treatment estimation with observational data, and further suggests a generic latent variable model for factorizing selection bias, as well as outcome. The author(s) build on this work and propose a set of deep generative models, with a hybrid objective function (generative + discriminative), that outperforms current approaches for ATE.