DR-VIDAL -- Doubly Robust Variational Information-theoretic Deep Adversarial Learning for Counterfactual Prediction and Treatment Effect Estimation on Real World Data

TL;DR

Observational data suffer from confounding biases that distort treatment-effect estimates. DR-VIDAL combines a VAE-based latent confounder decomposition, an Info-GAN counterfactual generator, and a doubly robust multitask predictor to estimate individualized treatment effects $\tau(\mathbf{x})$ from real-world data. Across synthetic benchmarks and real-world datasets IHDP, Twins, and Jobs, DR-VIDAL consistently outperforms TARNet, CEVAE, GANITE, and other baselines, with the doubly robust component contributing substantial gains. The approach is modular and open-source, offering a scalable tool for reliable counterfactual prediction and causal effect estimation in practical settings.

Abstract

Determining causal effects of interventions onto outcomes from real-world, observational (non-randomized) data, e.g., treatment repurposing using electronic health records, is challenging due to underlying bias. Causal deep learning has improved over traditional techniques for estimating individualized treatment effects (ITE). We present the Doubly Robust Variational Information-theoretic Deep Adversarial Learning (DR-VIDAL), a novel generative framework that combines two joint models of treatment and outcome, ensuring an unbiased ITE estimation even when one of the two is misspecified. DR-VIDAL integrates: (i) a variational autoencoder (VAE) to factorize confounders into latent variables according to causal assumptions; (ii) an information-theoretic generative adversarial network (Info-GAN) to generate counterfactuals; (iii) a doubly robust block incorporating treatment propensities for outcome predictions. On synthetic and real-world datasets (Infant Health and Development Program, Twin Birth Registry, and National Supported Work Program), DR-VIDAL achieves better performance than other non-generative and generative methods. In conclusion, DR-VIDAL uniquely fuses causal assumptions, VAE, Info-GAN, and doubly robustness into a comprehensive, performant framework. Code is available at: https://github.com/Shantanu48114860/DR-VIDAL-AMIA-22 under MIT license.

Figures (6)

• Figure 1: Directed acyclic graph modeling the causal relationships among treatment $t$, outcome $y$ and pre-treatment covariates $X$, under the latent space $Z$.
• Figure 2: Architecture of DR-VIDAL incorporating the variational autoencoder inferring the latent space (VAE), the generative adversarial network for calculating the counterfactual outcomes (GAN), and the doubly robust module (green box) for estimating ITE.
• Figure 3: Panel (a): performance (ATE) of DR-VIDAL vs. all other models on samples from the generative process of CEVAE. Panel (b) and (c): performance (PEHE) of DR-VIDAL with or without the doubly robust (DR, w/o DR) block vs. GANITE on samples from the generative process of CEVAE-GANITE.
• Figure 4: Performance comparison of doubly robust vs. non-doubly robust version of DR-VIDAL. The bar plots show how many times one model setup is better than the other in terms of error on the factual outcome ($y_f$). Panels, from left to right, show results on IHDP, Jobs and Twins datasets (100, 10, 100 iterations), respectively.
• Figure 5: Performance comparison of doubly robust vs. non-doubly robust version of DR-VIDAL. Panels, from left to right, show results on IHDP, Jobs and Twins datasets (100, 10, 100 iterations), respectively.