Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

An AI-powered Bayesian generative modeling approach for causal inference in observational studies

Qiao Liu, Wing Hung Wong

TL;DR

CausalBGM tackles causal inference in observational studies with high-dimensional covariates by learning a low-dimensional latent confounder space $Z=(Z_0,Z_1,Z_2,Z_3)$ within a fully Bayesian generative framework. It removes the encoder-decoder loop of prior AI-based methods, performing iterative mini-batch updates and modeling both mean and variance via Bayesian neural networks to yield well-calibrated posterior intervals for ADRF and ITE. The approach achieves superior or competitive accuracy across continuous and binary treatments, demonstrates reliable uncertainty quantification through calibrated posterior intervals, and scales to large datasets, underpinned by EGM-based initialization and a decoupled, parallelizable inference scheme. These contributions offer a principled, scalable, and interpretable tool for modern causal inference in genomics, healthcare, and social sciences, with public code and tutorials available.

Abstract

Causal inference in observational studies with high-dimensional covariates presents significant challenges. We introduce CausalBGM, an AI-powered Bayesian generative modeling approach that captures the causal relationship among covariates, treatment, and outcome. The core innovation is to estimate the individual treatment effect (ITE) by learning the individual-specific distribution of a low-dimensional latent feature set (e.g., latent confounders) that drives changes in both treatment and outcome. This individualized posterior representation yields estimates of the individual treatment effect (ITE) together with well-calibrated posterior intervals while mitigating confounding effect. CausalBGM is fitted through an iterative algorithm to update the model parameters and the latent features until convergence. This framework leverages the power of AI to capture complex dependencies among variables while adhering to the Bayesian principles. Extensive experiments demonstrate that CausalBGM consistently outperforms state-of-the-art methods, particularly in scenarios with high-dimensional covariates and large-scale datasets. By addressing key limitations of existing methods, CausalBGM emerges as a robust and promising framework for advancing causal inference in a wide range of modern applications. The code for CausalBGM is available at https://github.com/liuq-lab/bayesgm. The tutorial for CausalBGM is available at https://causalbgm.readthedocs.io.

An AI-powered Bayesian generative modeling approach for causal inference in observational studies

TL;DR

CausalBGM tackles causal inference in observational studies with high-dimensional covariates by learning a low-dimensional latent confounder space within a fully Bayesian generative framework. It removes the encoder-decoder loop of prior AI-based methods, performing iterative mini-batch updates and modeling both mean and variance via Bayesian neural networks to yield well-calibrated posterior intervals for ADRF and ITE. The approach achieves superior or competitive accuracy across continuous and binary treatments, demonstrates reliable uncertainty quantification through calibrated posterior intervals, and scales to large datasets, underpinned by EGM-based initialization and a decoupled, parallelizable inference scheme. These contributions offer a principled, scalable, and interpretable tool for modern causal inference in genomics, healthcare, and social sciences, with public code and tutorials available.

Abstract

Causal inference in observational studies with high-dimensional covariates presents significant challenges. We introduce CausalBGM, an AI-powered Bayesian generative modeling approach that captures the causal relationship among covariates, treatment, and outcome. The core innovation is to estimate the individual treatment effect (ITE) by learning the individual-specific distribution of a low-dimensional latent feature set (e.g., latent confounders) that drives changes in both treatment and outcome. This individualized posterior representation yields estimates of the individual treatment effect (ITE) together with well-calibrated posterior intervals while mitigating confounding effect. CausalBGM is fitted through an iterative algorithm to update the model parameters and the latent features until convergence. This framework leverages the power of AI to capture complex dependencies among variables while adhering to the Bayesian principles. Extensive experiments demonstrate that CausalBGM consistently outperforms state-of-the-art methods, particularly in scenarios with high-dimensional covariates and large-scale datasets. By addressing key limitations of existing methods, CausalBGM emerges as a robust and promising framework for advancing causal inference in a wide range of modern applications. The code for CausalBGM is available at https://github.com/liuq-lab/bayesgm. The tutorial for CausalBGM is available at https://causalbgm.readthedocs.io.
Paper Structure (20 sections, 16 equations, 3 figures, 3 tables)

This paper contains 20 sections, 16 equations, 3 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Methods
  3. Problem Setup
  4. Causal Generative Modeling
  5. Iterative Updating Algorithm
  6. Choice of latent dimension
  7. Model Initialization
  8. Model Hyperparameters
  9. Results
  10. Datasets
  11. Model Evaluation
  12. Baseline Methods
  13. Continuous Treatment Experiments
  14. Binary Treatment Experiments
  15. Posterior Interval
  16. ...and 5 more sections

Figures (3)

  • Figure 1: Illustration of CausalBGM framework. (A) The typical causal diagram in the observational study where the treatment, outcome, and covariates are observed variables. (B) The overview of CausalBGM model where variables are in rectangles and functions are in circles with incoming arrows indicating inputs to the function and outgoing arrows indicating outputs. $G$, $H$, and $F$ represent generative models for covariates, treatment, and outcome variables, respectively. $E$ represents the encoding function that creates circularity and is used for initialization purpose only. $E$ is removed in CausalBGM during the model training.
  • Figure 2: The performance of CausalBGM and baseline methods (Reg, DML with Lasso or neural network, and CausalEGM) under continuous treatment settings across three benchmark datasets. (A) Sun et al. dataset. (B) Hirano and Imbens dataset. (C) Colangelo and Lee dataset. The red curves represent the ground truth, while the blue curves indicate the estimated average dose-response of different methods with 95% confidence intervals based on 10 independent runs.
  • Figure 3: posterior interval analysis of CausalBGM using Imbens et al. dataset. (A) The calibration of empirical coverage rate at different treatment values ($x=0.2,0.5,1,1.5,2,2.5$). (B) The marginal density plot of treatment value $x$. Vertical dotted lines with different colors represent different treatment values ($x=0.2,0.5,1,1.5,2,2.5$) (C) The distribution of interval lengths at different significant levels $\alpha=0.01,0.05,0.1$. (D-F) The coverage indicator plots of CausalBGM at different significant levels $\alpha=0.01,0.05,0.1$ where the horizontal line indicates the truth average dose-response value at $x=2$, the “covered” intervals are marked in green, and “missed” intervals are marked in red.