Bayesian Additive Regression Trees for Propensity Scores

Why Use BART for Propensity Scores?

A propensity score is the probability of receiving treatment given measured covariates. Estimating it requires a model for treatment assignment, but the relationship between covariates and treatment may involve nonlinearities and interactions that are difficult to specify in advance. Bayesian additive regression trees (BART) provide a flexible way to represent such patterns: many small regression trees contribute to one fitted function, and Bayesian estimation accounts for uncertainty about the trees and their terminal-node values.

BART is not a guarantee of a good propensity model. It cannot account for confounders that were not measured, and its fitted probabilities do not by themselves demonstrate covariate balance or adequate overlap. Its role here is to provide a flexible estimate of treatment probabilities that must still be evaluated in the context of the intended causal analysis.

The BART Probit Model

For a binary treatment, one standard BART specification uses a probit link. Let \(m\) be the number of trees. Each tree \(j\) has a tree structure that partitions covariate space into terminal nodes, and a parameter \(\mu_{jk}\) for each terminal node \(k\). The tree returns the parameter of the terminal node into which the covariates fall. Summing the predictions from all trees gives the latent function \(f(X_i)\).

Here, \(\mathcal{T}_j\) denotes tree \(j\)’s structure, \(\mathcal{M}_j\) its terminal-node parameters, and \(g\) the prediction from that tree. The standard normal cumulative distribution function \(\Phi\) maps the real-valued sum to a probability between zero and one. A single tree makes a piecewise-constant prediction; adding trees allows the model to build more varied patterns, including nonlinearities and interactions, while keeping each tree relatively simple.

The probit specification has a useful latent-variable representation. Conditional on the covariates and the tree ensemble, define \(Z_i=f(X_i)+\varepsilon_i\), where \(\varepsilon_i\sim N(0,1)\), and set \(T_i=1\) exactly when \(Z_i>0\). Then

The error variance is fixed at one to set the probit scale; it is not a freely estimated residual variance in this binary-treatment model. In the usual likelihood, treatment indicators are modeled as conditionally independent given their covariates and the model parameters. The covariates used for estimating treatment assignment should be measured before treatment, not variables that may have been affected by treatment.

How the Bayesian Tree Sum Is Regularized

The flexibility of a large tree ensemble needs regularization. A common BART prior makes deep trees less likely than shallow ones. For a node at depth \(d\), the probability of splitting it is often set to \(\alpha/(1+d)^\beta\), with parameters chosen so that splitting becomes less probable as depth increases. Candidate split variables and cut points are selected according to specified prior rules. Terminal-node values are given a centered prior, commonly a normal distribution with a scale small enough that any one tree is encouraged to make a modest contribution.

The point of these priors is not to assert that the true treatment mechanism consists of a particular number of tiny trees. Rather, the model combines many regularized components so that the sum can express a richer function while limiting the influence of any one component. Exact prior choices and sampling details vary among implementations, so an analysis should report the BART specification and software settings used.

Bayesian computation produces draws of the tree structures and their terminal-node values from the posterior distribution, conditional on the observed treatment indicators and covariates. For a new covariate pattern \(x\), each retained draw gives a latent prediction \(f^{(s)}(x)\), and thus a draw-specific probability \(\Phi(f^{(s)}(x))\). The usual point estimate of the propensity is the posterior mean of those probabilities:

where \(S\) is the number of retained posterior draws. The transformation is applied within each draw before averaging. In general, the average of transformed values is not the same as transforming the average latent prediction.

Posterior Predictions: Average After Applying the Link

Because the tree structures and terminal-node values vary across posterior draws, the latent prediction can vary too. The point prediction for a covariate pattern summarizes this variation by averaging the draw-specific treatment probabilities. The following example shows why the order of operations matters.

How Tree Sums Represent Different Covariate Patterns

A tree can split on one covariate and then split again on another within one branch. Such a structure can represent an interaction: the contribution of one covariate may depend on the value of another. Across trees, different splits can combine into a flexible function without requiring the analyst to list every possible interaction in advance.

Checking the Estimate and Understanding Its Limits

A flexible fitted model still needs diagnostics. Begin by checking that the estimated probabilities are finite, within zero and one, and plausible for the observed treatment patterns. Inspect whether the model assigns probabilities extremely close to zero or one; such estimates can signal limited empirical overlap or a highly confident fit in regions with little data. A predicted probability alone cannot distinguish these explanations.

Evaluate probability predictions with suitable predictive checks, such as calibration summaries, and use the cross-validation principles already introduced when comparing model settings. Predictive classification accuracy alone is not enough: a model can classify treatment status well while giving probabilities that are poorly calibrated or unsuitable for the intended analysis. After using the estimated scores in a causal design, assess covariate balance and overlap directly. A fitted propensity score is not itself evidence that either condition has been achieved.

BART depends on its prior choices, the covariates supplied, the sample information in each region of covariate space, and the adequacy of the binary-outcome model. Sparse regions may produce unstable or prior-sensitive predictions. Poorly chosen covariates, post-treatment variables, measurement error, and unmeasured confounding remain problems that tree flexibility cannot solve. BART also does not automatically solve a lack of overlap: if treated and control units have very different covariate patterns, a flexible fit cannot create data where none exist.

Its advantages and trade-offs are clearer when compared with common alternatives:

BART’s tree sum is an ensemble within one Bayesian model. It is different from combining distinct propensity estimators, a separate strategy. Nor does Bayesian estimation mean that the model is automatically well specified: posterior predictions are conditional on the data, likelihood, and prior used.

Common Mistakes

• Applying the link to only one latent prediction. A value from one posterior draw is not the posterior-mean propensity; average \(\Phi(f^{(s)}(x))\) across draws.
• Treating the tree sum as a probability. The sum \(f(x)\) can be any real number. Apply \(\Phi\) to obtain a probability in the probit model.
• Calling a fitted score a causal effect. It predicts treatment assignment conditional on measured covariates; it does not quantify the effect of treatment.
• Assuming flexibility guarantees balance. Balance and overlap must be assessed in the intended analysis rather than inferred from the model class.
• Using information measured after treatment. Such variables may be affected by treatment and can make a treatment-assignment model inappropriate for the intended causal design.