Bayesian Decision Trees for Confounder Selection in Mediation Analysis

ICLR 2026 Conference Submission14565 Authors

Keywords: Causal mediation, Bayesian additive regression trees, Dirichlet prior

TL;DR: We develop a Bayesian nonparametric framework for data-driven confounder selection in mediation analysis, designed to handle high-dimensional settings with numerous potential confounders.

Abstract: The estimation of causal effects in observational research fundamentally relies on proper adjustment for confounding variables. As such, identifying relevant confounders from the data is an important preliminary task. Although numerous data-driven techniques have been suggested for confounder selection in conventional exposure-outcome analyses, such methodologies are absent in causal mediation analysis, which entails identifying the effects within the exposure-mediator-outcome framework. This paper presents a Bayesian framework for confounder selection in mediation analysis via Bayesian additive regression trees (BART). We specify separate models for the exposure, mediator, and outcome, and introduce a common sparsity-inducing prior on their selection probability vectors. This enables identification of covariates that are jointly important across all models (that is, potential confounders). Furthermore, we introduce a novel criterion for confounder selection in the context of mediation analysis and establish that satisfaction of this criterion ensures the validity of the sequential ignorability assumption under certain conditions. The proposed method demonstrates consistently strong performance across a range of simulation scenarios, offering a practical approach for confounder selection in high-dimensional mediation analysis.

Primary Area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)

Submission Number: 14565