Direct Bayesian Additive Regression Trees for Conditional Average Treatment Effects in Regression Discontinuity Designs
| Authors | Daisuke Kondo & Shonosuke Sugasawa |
| Year | 2026 |
| Field | AI / ML |
| arXiv | 2603.03819 |
| Download | |
| Categories | stat.ME, stat.ML |
Abstract
Regression discontinuity designs (RDD) are widely used for causal inference. In many empirical applications, treatment effects vary substantially with covariates, and ignoring such heterogeneity can lead to misleading conclusions, which motivates flexible modeling of heterogeneous treatment effects in RDD. To this end, we propose a Bayesian nonparametric approach to estimating heterogeneous treatment effects based on Bayesian Additive Regression Trees (BART). The key feature of our method lies in adopting a general Bayesian framework using a pseudo-model defined through a loss function for fitting local linear models around the cutoff, which gives direct modeling of heterogeneous treatment effects by BART. Optimal selection of the bandwidth parameter for the local model is implemented using the Hyvärinen score. Through numerical experiments, we demonstrate that the proposed approach flexibly captures complicated structures of heterogeneous treatment effects as a function of covariates.
Engineering Breakdown
Plain English
This paper addresses a critical problem in causal inference: when treatment effects vary across different groups or conditions, standard regression discontinuity designs (RDD) can produce misleading average estimates. The authors propose a Bayesian nonparametric method that combines BART (Bayesian Additive Regression Trees) with a pseudo-model framework to directly estimate heterogeneous treatment effects around a cutoff threshold. The key innovation is fitting local linear models at the discontinuity boundary while letting BART flexibly capture how treatment effects vary with covariates, automatically selecting the bandwidth parameter for optimal local fitting. This approach allows practitioners to understand not just whether a treatment works, but for whom and how much.
Core Technical Contribution
The core novelty lies in the integration of BART with a Bayesian pseudo-model framework specifically designed for RDD settings. Instead of estimating a single average treatment effect (ATE), the method directly models heterogeneous treatment effects (HTE) by building trees that respect the local linear regression structure required near the discontinuity cutoff. The pseudo-model approach using loss functions enables seamless incorporation of RDD constraints (local linearity, bandwidth selection) into the Bayesian nonparametric framework, whereas prior BART applications to causal inference required post-hoc adjustments or worked in less structured settings. This is the first work to systematically combine flexible tree-based heterogeneity modeling with the technical requirements of discontinuity designs.
How It Works
The method operates in several stages: first, data near the cutoff threshold is partitioned based on treatment assignment (above/below cutoff). A pseudo-likelihood function is constructed based on local linear regression losses—this creates a Bayesian objective that respects RDD's requirement for linearity in small neighborhoods around the cutoff. BART then builds an ensemble of regression trees on top of this pseudo-likelihood, where each tree split represents a decision rule that partitions the covariate space, and the final prediction for each leaf node represents a local treatment effect estimate. The bandwidth parameter (controlling how much data to use locally) is selected automatically via cross-validation or marginal likelihood optimization integrated into the Bayesian workflow. The output is a posterior distribution over heterogeneous treatment effects, giving both point estimates and uncertainty quantification for each covariate subgroup.
Production Impact
For practitioners, this method directly addresses a major limitation in A/B testing and policy evaluation: understanding when and for whom treatments actually matter. In production settings like healthcare, education, or targeted marketing, practitioners often discover that average treatment effects hide critical subgroup dynamics—some populations benefit greatly while others don't, and standard RDD reports only the average. Adopting this approach means your causal inference pipeline can automatically discover and quantify these effect modifications without manual subgroup pre-specification, reducing the risk of rolling out one-size-fits-all policies. However, the trade-offs are real: BART requires substantial tuning (priors, tree depth, number of trees), the pseudo-model framework adds computational overhead compared to standard local linear regression, and the method assumes the RDD assumptions (local continuity, no self-selection at cutoff) hold—which needs careful validation. For moderate-scale datasets (thousands to tens of thousands of samples), this is implementable, but scaling to millions of rows requires careful engineering of the tree growing and bandwidth selection steps.
Limitations and When Not to Use This
The paper assumes strict RDD validity: that units cannot precisely manipulate their running variable to land on their preferred side of the cutoff, and that potential outcomes are continuous at the boundary. In practice, strategic sorting (e.g., students trying to stay just below an admission score threshold) violates these assumptions and isn't addressed here. The method also inherits BART's interpretability limitations—while you get heterogeneous effect estimates, understanding why effects differ across groups requires post-hoc interpretation of tree splits, which can be opaque. Bandwidth selection is mentioned but not fully detailed in the abstract; if this relies on standard cross-validation, it may overfit in small samples or high-dimensional covariate spaces. Finally, the computational cost of both BART and bandwidth optimization is likely substantial compared to simple local linear regression, but no runtime comparisons or scalability analysis is mentioned, leaving uncertainty about practical feasibility for large-scale applications.
Research Context
This work extends a rich literature on heterogeneous treatment effects estimation (building on Athey & Wager, Kunzel et al., and others who applied BART to causal forests) into the structured setting of regression discontinuity designs. RDD is a gold-standard method in applied microeconomics and policy evaluation because it exploits natural cutoffs (test scores determining school placement, income thresholds for benefits eligibility) to approximate randomization; the paper acknowledges that many RDD applications show substantial effect heterogeneity but lack principled tools to model it. By combining BART's flexibility with RDD's structure, the authors open a direction for other local/structured causal inference methods (propensity score matching, instrumental variables) to similarly incorporate heterogeneous effect modeling. This likely positions future work on combining tree-based and kernel-based methods with other causal designs.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::