Predictive State Propensity Subclassification (PSPS): A causal inference algorithm for data-driven propensity score stratification

Joseph Kelly, Jing Kong, Georg M. Goerg

Keywords: causal inference, multi-task learning, predictive state representation, propensity score matching, causal representation learning

TL;DR: Predictive State Propensity Subclassification is a novel causal inference algorithm built on multi-outcome networks, predictive state representations, and proper causal loss functions.

Abstract: We introduce Predictive State Propensity Subclassification (PSPS), a novel learning algorithm for causal inference from observational data. PSPS combines propensity and outcome models into one encompassing probabilistic framework, which can be jointly estimated using maximum likelihood or Bayesian inference. The methodology applies to both discrete and continuous treatments and can estimate unit-level and population-level average treatment effects. We describe the neural network architecture and its TensorFlow implementation for likelihood optimization. Finally we demonstrate via large-scale simulations that PSPS outperforms state-of-the-art algorithms – both on bias for average treatment effects (ATEs) and RMSE for unit-level treatment effects (UTEs).