Google Research

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

Google LLC (2020)


We introduce Predictive State Propensity Subclassification (PSPS), a novel estimation method for undertaking causal inference from observational studies. The methodology applies to both discrete and continuous treatments and can estimate unit-level and population-level average treatment effects. PSPS combines propensity and outcome models into one encompassing probabilistic model, which can be jointly estimated using maximum likelihood or Bayesian inference. We give a detailed overview on the TensorFlow implementation for likelihood optimization and show via large-scale simulations that it outperforms several state of the art methods -- both in terms of bias and variance for average as well as unit-level treatment effects. Finally we illustrate the methodology and algorithms on standard datasets in the literature.

Learn more about how we do research

We maintain a portfolio of research projects, providing individuals and teams the freedom to emphasize specific types of work