A Bayesian Additive Regression Tree Model for Learning Conditional Average Treatment Effects in Regression Discontinuity Designs

arXiv:2503.00326 · stat.ME, stat.ML · Submitted 2025-03-01 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: We're looking at the paper: "A Bayesian Additive Regression Tree Model for Learning Conditional Average Treatment Effects in Regression Discontinuity Designs".

Jane: Here's the gist of the paper. The gist This paper develops an effective Bayesian approach to conditional average treatment effect (CATE) estimation in regression discontinuity designs (RDD), which is an increasingly prevalent form of quasi-experiment that facilitates causal inference.

Tom: Model Development The new model is a variant of a Bayesian additive regression tree (BART) model with linear leaflevel regressions on the running variable and a treatment dummy (and their interaction).

Jane: This approach adaptively partitions covariate space into regions where the slope on the running variable appreciably differs, providing interpretable Bayesian inference on conditional average treatment effects near the cutoff.

Tom: The justification for developing this novel and non-trivial extension to BART is simply this: BART is widely-acknowledged to be superior at estimating conditional expectations compared to the methods that underpin previous RDD approaches to CATE estimation (linear models with pre-specified nonlinear basis functions or single individual regression tree models).

Paper summary: Tom: So looking at the title, "A Bayesian Additive Regression Tree Model for Learning Conditional Average Treatment Effects in Regression Discontinuity Designs," it summarizes the whole effort to use this specific tree ensemble method for causal effect estimation in RDDs.

Jane: And what it means simply is that they found a way to estimate how much treatment effects change depending on other variables using data-driven trees instead of rigid mathematical assumptions ><ref:2503.00326#pg1>.

Lu: It’s about solving the problem of needing regression adjustment for causal inference where you can't just use global splines, and this paper introduces a Bayesian tree ensemble that does it without assuming a pre-specified basis expansion ><ref:2503.00326#pg2>.

Meng: For someone shipping models to production, the implication is that we might be able to deploy these types of ensemble methods for more complex CATE modeling where we can't just rely on simple linear assumptions about the underlying functions ><ref:2503.00326#pg1>.

Lalam: I'm pretty sure this work settles a notable flaw in the existing literature by showing that data-driven CATE discovery is possible on realistically large data sets ><ref:2503.00326#pg1>.

Tom: It moves the field by providing a concrete, non-parametric approach to handling conditional average treatment effects in regression discontinuity designs, which is a very specific and important causal inference scenario ><ref:2503.00326#pg1>.

Jane: So that’s what this paper does—it proposes a model that fits the complexity of the data rather than forcing the data into a pre-set mathematical box ><ref:2503.00326#pg1>.

Conclusion: Tom: So we've been looking at this paper on Bayesian Additive Regression Trees for estimating conditional average treatment effects in regression discontinuity designs, or RDDs, and what it actually does.

Jane: Right, so to recap, they put together a new model using these tree structures to figure out how the treatment effect changes based on other variables near that cutoff point where the intervention happens.

Lu: Exactly. They’re taking something like regression discontinuity—which is a way to do causal inference when you have data that looks like it comes from an experiment, but isn't—and they are using these Bayesian additive regression trees to estimate those conditional averages.

Meng: Okay, so the main point I see for deploying this is that instead of picking some fixed mathematical formula beforehand, which is what older methods usually do, this model lets the data guide the structure of how we look at things.

Tom: That’s right. And they found that by using these BART models—Bayesian Additive Regression Trees—they can partition the covariate space into different regions where things are behaving differently near the cutoff, and that gives us a really interpretable Bayesian inference on those effects.

Jane: Interpretable is a big word. So what does that actually mean for someone listening who isn't deep in statistics? It means they can see *why* the treatment effect is different in one area compared to another, instead of just getting one single number for the whole thing.

Lu: It lets you see if course load, for instance, might be a big factor that changes how that treatment effect behaves right at the point of cut-off. The paper shows this tree ensemble approach is generally better than some of the previous methods they compared it against.

Meng: Better how? Because we're talking about real data here, not some clean simulation. What matters to me is that when you try to run this on a massive dataset, it doesn’t just break down or give you garbage results due to needing too much pre-specification.

Tom: The authors show in their simulations that the BARDDT model produces lower root-mean-squared estimation error than other estimators they tested, and it tends to beat the simpler homogeneous estimator most of the time. It seems robust for large datasets.

Jane: So for a researcher trying to understand causal effects in real-world scenarios, this suggests that we don't always need those overly complex, pre-defined mathematical functions if we use these flexible AI models instead.

Lu: It moves us toward using data-driven discovery rather than just fitting a known shape to the data. That’s the direction they are pushing with this model extension of BART.

Meng: And for me, it means that if we can build software around this, it could be useful for situations where we know the effect is nonlinear but we don't want to spend all day guessing the exact functional form upfront.

Tom: So to wrap up, this paper presents a way to use these tree ensembles to do conditional average treatment effect estimation in RDDs without needing that steep model specification requirement that trips up other frequentist methods.

Jane: It’s moving from assuming a formula is correct, to letting the data tell us where the effect is changing. That’s what this paper aims to achieve with their Bayesian Additive Regression Tree model.

Lu: The next thing we need to look at is how these results translate when you try to apply them in a different kind of causal setup entirely.

Rafael Alcantara, P. Richard Hahn, Hedibert F. Lopes

University of Texas at Austin · Arizona State University · Insper

stat.ME, stat.ML

Submitted: 2025-03-01

Updated: 2026-10-16

Code: https://github.com/rafaelcalcantara/BART-RDD

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 83/100

The gist: The gist This paper develops an effective Bayesian approach to conditional average treatment effect (CATE) estimation in regression discontinuity designs (RDD), which is an increasingly prevalent

Key concepts

Regression Discontinuity Design (RDD): RDD is a quasi-experiment where treatment assignment is determined by whether an observed variable crosses a specific cutoff point. It allows researchers to estimate causal effects by comparing outcomes on either side of this sharp threshold, making it useful for causal inference.
Conditional Average Treatment Effect (CATE): CATE refers to the average treatment effect for a specific individual or group, given their observed characteristics (covariates). Instead of estimating one overall effect, CATE estimates how the treatment effect varies across different subgroups within the population.
Bayesian Additive Regression Tree (BART) Model: BART is a powerful machine learning technique used here as a flexible model. It builds predictions by combining many simple decision trees. This ensemble approach allows the model to capture complex, non-linear relationships in the data without needing to pre-specify the exact functional form.
BARDDT Model Structure: The BARDDT model uses a piecewise linear function within the BART ensemble. It calculates CATE at specific points by summing contributions from individual trees, using priors on coefficients to regularize and constrain the estimated treatment effects.

Terminology

Bayesian Inference — 1× in this paper · explained in 6 episodes
Causal Inference — 1× in this paper · explained in 4 episodes

Summary

The gist This paper develops an effective Bayesian approach to conditional average treatment effect (CATE) estimation in regression discontinuity designs (RDD), which is an increasingly prevalent form of quasi-experiment that facilitates causal inference.

Model Development

The new model is a variant of a Bayesian additive regression tree (BART) model with linear leaflevel regressions on the running variable and a treatment dummy (and their interaction). This approach adaptively partitions covariate space into regions where the slope on the running variable appreciably differs, providing interpretable Bayesian inference on conditional average treatment effects near the cutoff. The justification for developing this novel and non-trivial extension to BART is simply this: BART is widely-acknowledged to be superior at estimating conditional expectations compared to the methods that underpin previous RDD approaches to CATE estimation (linear models with pre-specified nonlinear basis functions or single individual regression tree models).

Identification in RDDs

In the context of regression discontinuity designs, treatment effects at the cutoff can be estimated as the magnitude of a discontinuity in the conditional mean response function at that point. When incorporating additional covariates, CATE estimation boils down to estimation of conditional expectation functions E[Y | X = x, W = w, Z = z], for which BART models are used.

BARDDT Model Structure

The BARDDT model uses a piecewise linear function within a BART ensemble. The prediction function for tree j is defined as gj(x, w, z) = ψ(x, z)Γbj (x,w). This leads to the CATE at x = 0 being a sum of the ∆bj(0,w) elements across all trees j. The priors on the ∆ coefficients directly regularize the treatment effect.

Simulation and Evaluation

The paper describes a parametrized protocol for simulating data for evaluating CATE estimation methods in RDDs. The evaluation criteria used is the average root-mean-squared estimation error, expressed as a fraction of a default ATE estimator. The results show that BARDDT produces lower RMSE than the other estimators on average and is most often better than the homogeneous estimator.

Conclusion

BARDDT solves all of these problems, permitting data-driven CATE discovery on realistically large data sets. The software for fitting BARDDT is freely available in the stochtree package, available in both R and Python. This approach brings to bear an effective modern computational strategy (tree ensembles) for nonlinear CATE estimation, rather than merely assuming the problem away. The proposed family of DGPs addresses a notable flaw in the existing literature. This protocol allows for a much richer exploration of estimator properties that is not tied to a particular modelling assumption and that can be easily adapted to study estimator properties in more specific use cases. The resulting effect moderation tree indicates that course load (credits attempted) in the academic term leading to their probation is a strong moderator. The authors do not explore course load and age at all as potential moderators, which our results strongly suggest. The correlation between our CATE predictions and each of the null features is effectively zero. This demonstrates the ability of BARDDT to handle more features. The paper demonstrates that BARDDT estimates RDD CATEs successfully and markedly better than available alternatives.

--- Page 1 ---

A Bayesian Additive Regression Tree Model for Learning Conditional Average Treatment Effects in Regression Discontinuity Designs Rafael Alcantara∗1, P. Richard Hahn2, and Hedibert F. Lopes3 The University of Texas at Austin Arizona State University Insper July 28, 2026 Abstract This paper develops an effective Bayesian approach to conditional average treatment effect (CATE) estimation in regression discontinuity designs (RDD), which is an increasingly prevalent form of quasi-experiment that facilitates causal inference. Earlier Bayesian approaches do not easily accommodate CATE estimation while recent frequentist approaches to this problem assume a known basis expansion, a steep model specification requirement that our approach avoids. The new model is a variant of a Bayesian additive regression tree (BART) model with linear leaflevel regressions on the running variable and a treatment dummy (and their interaction).

Improvements for AI systems

Stronger CATE Estimation in RDDs via BARDDT Model Development: The improved AI system can estimate CATE functions at x = 0; i.e., τ (0, w), for some covariate vector w by using a variant of BART with linear leaf-level regressions on the running variable and a treatment dummy (and their interaction), which adaptively partitions the covariate space into regions where the slope on the running variable appreciably differs.
Robustness Against Model Specification: The system can avoid the steep model specification requirement of frequentist approaches by using a Bayesian additive regression tree (BART) model, which is widely-acknowledged to be superior at estimating conditional expectations compared to the methods that underpin previous RDD approaches.
Causal Inference via Explicit Regularization: The system can estimate the treatment effect function by parameterizing it as a sum of regression trees where the leaf node parameters are a vector of regression coefficients, allowing for explicit regularization through priors on the ∆ coefficients, which is noted to be a way to regularize the treatment effect.
Adaptive Local Bandwidth Determination: The system can determine when to prune regions away from the cutoff value because, in BARDDT, different trees in the ensemble are effectively using different local bandwidths and these fits are then blended together, which is described as a way to avoid needing a 'bandwidth' as is used in the polynomial approach to RDD.
Generalizability Across DGPs: The system can perform CATE estimation across a variety of data generating processes by accommodating complex response surfaces while allowing direct control over the variation and magnitude of the implied CATEs, moving beyond overly-simplistic response surfaces with improbably large treatment effects and implausibly conspicuous heterogeneity.