Bayesian Nonparametrics for Principal Stratification with Continuous Post-Treatment Variables

2024-05-27

AI Summary

This paper addresses the challenge of principal stratification causal identification under continuous post-treatment mediators. We propose a Bayesian nonparametric principal stratification method based on a confounders-aware shared-atom mixture model coupled with a hierarchical Dirichlet process prior. The approach enables data-adaptive coarse-graining of latent principal strata, interpretable definition of principal effects, and quantification of individual stratum membership uncertainty. In simulations, it substantially improves identification accuracy and robustness over existing methods. Applied to U.S. air quality regulation policy evaluation, it precisely estimates stratified causal effects of pollution exposure on health outcomes. Our key innovation is the first integration of a confounder-aware mechanism into a shared-atom nonparametric modeling framework—simultaneously resolving three core challenges: principal stratum identification, effect interpretability, and uncertainty characterization.

Technology Category

• Reasoning under Uncertainty: Causality
• Machine Learning: Calibration & Uncertainty Quantification
• Intelligent Robots: State Estimation

Application Category

• User Modeling, Personalization and Recommendation: Explainable and interpretable methods for personalization
• Graph Algorithms and Modeling for the Web: Algorithms and analysis for incomplete, noisy, or partially observed Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Principal stratification provides a causal inference framework for investigating treatment effects in the presence of a post-treatment variable. Principal strata play a key role in characterizing the treatment effect by identifying groups of units with the same or similar values for the potential post-treatment variable under both treatment levels. The literature has focused mainly on binary post-treatment variables. Few papers considered continuous post-treatment variables. In the presence of a continuous post-treatment, a challenge is how to identify and characterize meaningful coarsening of the latent principal strata that lead to interpretable principal causal effects. This paper introduces the confounders-aware shared-atom Bayesian mixture, a novel approach for principal stratification with binary treatment and continuous post-treatment variables. Our method leverages Bayesian nonparametric priors with an innovative hierarchical structure for the potential post-treatment variable that overcomes some of the limitations of previous works. Specifically, the novel features of our method allow for (i) identifying coarsened principal strata through a data-adaptive approach and (ii) providing a comprehensive quantification of the uncertainty surrounding stratum membership. Through Monte Carlo simulations, we show that the proposed methodology performs better than existing methods in characterizing the principal strata and estimating principal effects of the treatment. Finally, our proposed model is applied to a case study in which we estimate the causal effects of US national air quality regulations on pollution levels and health outcomes.

Problem

• Develops Bayesian nonparametric method for continuous post-treatment variables
• Identifies interpretable principal strata for causal effects
• Improves uncertainty quantification in stratum membership

Innovation

• Bayesian nonparametric priors for continuous post-treatment variables
• Confounders-aware shared-atom Bayesian mixture model
• Data-adaptive coarsening of principal strata