An End-to-end Learning Approach for Counterfactual Generation and Individual Treatment Effect Estimation

Estimating the causal effect due to an intervention is important for many applications, such as healthcare. Unobserved counterfactuals make unbiased treatment effect estimation non-trivial. Among existing approaches, counterfactual generation which augments observational data with generated pseudo counterfactuals has been found promising for reducing the bias. These methods typically take a two-stage approach for the counterfactual generation and treatment effect estimation. Therefore, the counterfactual generation could be sub-optimal. To this end, we propose to jointly optimize the auxiliary models for generating the counterfactuals and the outcome estimation models. In particular, we demonstrate the viability by first connecting a counterfactual outcome generator with a reparameterized VAE model, and then learning them in an end-to-end fashion using the EM algorithm. Our evaluation results based on synthetic and semi-synthetic datasets show that a simple causal effect VAE model learned together with the counterfactual outcome generator can outperform a number of SOTA models for treatment effect estimation.