Bayesian regression tree models for causal inference: regularization, confounding, and heterogeneous effects

P. Richard Hahn; Jared S. Murray; Carlos Carvalho

Source-linked AI summary

Bayesian regression tree models for causal inference: regularization, confounding, and heterogeneous effects

P. Richard Hahn, Jared S. Murray, Carlos Carvalho

arXiv:1706.09523v4Jun 29, 2017stat.ME

TL;DR

The paper addresses biased and insufficiently controlled heterogeneous-treatment-effect estimates from flexible models under strong confounding and small effects. It proposes Bayesian causal forests that incorporate an estimated propensity score and separately regularize prognostic and treatment-effect functions, performing especially well in strongly confounded, targeted-selection, weak-effect settings. An observational smoking study illustrates the resulting heterogeneous-effect inferences.

Problem
Flexible predictive models can produce biased treatment-effect estimates under strong confounding and inadequate control of heterogeneity regularization when prognostic and treatment-effect variation differ.
Method
Bayesian causal forests use separate BART forests for prognostic and treatment effects and incorporate an estimated propensity score into the response-surface prior.
Results
Simulations show Bayesian causal forests perform at least as well as existing approaches across plausible data-generating processes and dramatically better in many strongly confounded, targeted-selection, weak-effect settings.
Takeaways & Limitations
The model supports informative shrinkage toward homogeneous effects while preserving flexible inference on heterogeneous treatment effects in observational data.
Takeaways & Limitations
The propensity-score adjustment works only when the estimated propensity score accurately approximates the true propensity score, motivating sensitivity analysis across estimates.

Abstract

from arXiv · show

This paper presents a novel nonlinear regression model for estimating heterogeneous treatment effects from observational data, geared specifically towards situations with small effect sizes, heterogeneous effects, and strong confounding. Standard nonlinear regression models, which may work quite well for prediction, have two notable weaknesses when used to estimate heterogeneous treatment effects. First, they can yield badly biased estimates of treatment effects when fit to data with strong confounding. The Bayesian causal forest model presented in this paper avoids this problem by directly incorporating an estimate of the propensity function in the specification of the response model, implicitly inducing a covariate-dependent prior on the regression function. Second, standard approaches to response surface modeling do not provide adequate control over the strength of regularization over effect heterogeneity. The Bayesian causal forest model permits treatment effect heterogeneity to be regularized separately from the prognostic effect of control variables, making it possible to informatively "shrink to homogeneity". We illustrate these benefits via the reanalysis of an observational study assessing the causal effects of smoking on medical expenditures as well as extensive simulation studies.

1 Introduction

The paper develops Bayesian causal forests for heterogeneous treatment-effect estimation under small effects, strong confounding, and distinct prognostic and treatment-effect variation. Simulations and an observational smoking study motivate separate regularization and propensity-score adjustment.

1 Introduction: Regularization designed for prediction can bias causal estimates under confounding and inadequately control conditional treatment-effect estimator variance.These problems become especially important when prognostic variation differs substantially from treatment-effect heterogeneity.
1 Introduction: Bayesian causal forests model the response with a flexible sum of regression trees and incorporate an estimated propensity score into the response-surface prior.The propensity-score transformation is intended to improve treatment-effect estimation under moderate to strong confounding, particularly targeted selection.
1 Introduction: The model decomposes the response into separate forests for prognostic effects and treatment effects, allowing treatment-effect heterogeneity to be regularized independently.Under this parameterization, standard regression-tree priors shrink treatment effects toward homogeneity.
1 Introduction: The parameterization interpolates between separate treated/control models and treating assignment as another covariate, making the prior over treatment effects more interpretable and controllable.Separate models provide no borrowing or regularization, whereas the covariate approach can be difficult to interpret with flexible models.
1 Introduction: Simulations find Bayesian causal forests at least as effective as existing approaches across plausible data-generating processes and dramatically better in many strongly confounded, targeted-selection, weak-effect settings.The paper also reanalyzes smoking effects on medical expenditures using average and conditional average treatment-effect inferences.
1 Introduction: The paper targets a regime of small effect sizes, heterogeneous effects, and strong confounding that prior research did not specifically emphasize despite its applied importance.The broader literature includes nonlinear regression, tree ensembles, and methods for average or conditional treatment effects.

2 Problem statement and notation

The paper frames treatment-effect estimation using observed outcomes, binary treatment, and control variables under strong ignorability. Its contribution is a prior specification that models prognostic and treatment effects separately while incorporating the propensity score.

2 Problem statement and notation: The setup uses a scalar response Y, binary treatment Z, d-dimensional controls x, and n independent observations (Y_i, Z_i, x_i).Capital letters denote random variables, lowercase letters realized values, and X is the n×d control matrix.
2 Problem statement and notation: CATE is the difference in responses between treatment settings Z_i=1 and Z_i=0, averaged over subpopulations defined by covariates x.The estimand is expressed through potential outcomes in the potential-outcomes framework.
2 Problem statement and notation: The analysis assumes strong ignorability: no unmeasured confounders and overlap across covariate space are required for treatment-effect estimation.Under these conditions, potential-outcome means equal observed conditional means by treatment level.
2 Problem statement and notation: The proposed prior places independent BART priors on prognostic function μ and treatment-effect function τ, using an estimated propensity score π̂(x_i).The contribution is a prior specification intended to improve estimation of τ rather than merely prediction of f.

3 Bayesian additive regression trees for heterogeneous treatment effect estimation

BART represents an unknown function as a sum of weak, piecewise-constant binary regression trees. For treatment-effect estimation, a direct prior on the response surface induces a difficult-to-control prior on heterogeneity.

3 Bayesian additive regression trees for heterogeneous treatment effect estimation: BART is motivated by prior evidence of strong predictive performance and suitability for interactions, discontinuities, monotone covariate transformations, and limited tuning.The paper extends this predictive modeling framework to heterogeneous treatment-effect estimation.
3 Bayesian additive regression trees for heterogeneous treatment effect estimation: BART expresses f(x) as a sum of many piecewise-constant binary regression trees whose leaves assign parameters to covariate-space partitions.Each tree contains internal splitting rules and terminal-node parameters defining a step function.
3 Bayesian additive regression trees for heterogeneous treatment effect estimation: The forest prior favors weak learners through small trees and leaf parameters near zero, assigning extremely low prior probability to large, deep trees.The splitting probability decreases with node depth, using η=0.95 and β=2 in the cited specification.
3 Bayesian additive regression trees for heterogeneous treatment effect estimation: The induced marginal prior for f(x) is centered at zero, with approximately 95% of pointwise prior mass within ±2σ_0.σ_0 calibrates the plausible range of the regression function.
3 Bayesian additive regression trees for heterogeneous treatment effect estimation: A direct BART prior on f(x,z) makes the induced prior on τ(x)=f(x,1)−f(x,0) depend on covariate dimension and dependence with treatment.This makes regularization of heterogeneous effects difficult to understand and control.

4 The central role of the propensity score in regularized causal modeling

The paper identifies regularization-induced confounding as a source of biased heterogeneous-effect estimates. Including an estimated propensity score as a covariate reduces finite-sample bias in regularized treatment-effect estimation.

4 The central role of the propensity score in regularized causal modeling: Including an estimated propensity score as a covariate reduces bias in regularized treatment-effect estimates in finite samples.The authors recommend this adjustment as routine practice across models or algorithms used for heterogeneous-effect estimation.
4 The central role of the propensity score in regularized causal modeling: Regularization is necessary for nonparametric or high-dimensional heterogeneous-effect estimation, but it can interact with confounding through regularization-induced confounding and targeted selection.These concepts motivate the propensity-score adjustment developed in the paper.

4.1 Regularization-induced confounding

Regularization can distort treatment-effect estimation even when treatment effects are identified by the conditional expectation function. In finite samples, the prior over the response surface can substantially influence posterior treatment-effect estimates.

Likelihood-based identification of treatment effects does not resolve finite-sample estimation concerns because many response functions can fit the data similarly while implying different effects.
Posterior treatment-effect estimates can be substantially influenced by the prior distribution over the response surface at realistic sample sizes.The paper calls this phenomenon regularization-induced confounding (RIC).

Example: RIC in the linear model

The linear example shows how regularization-induced confounding arises when a regularized prior favors prognostic coefficients that can absorb treatment-related structure. The resulting bias depends on confounding, treatment selection, and unknown nuisance parameters rather than being directly controlled by the analyst.

In the linear model, the Bayes estimator is the posterior mean of the treatment effect under a Gaussian prior on regression coefficients.The analysis examines the bias of the posterior mean under a homogeneous treatment effect and linear response and treatment models.
A ridge prior regularizes control-variable coefficients while leaving the treatment-effect prior flat, creating unequal prior treatment across model components.The ridge parameter is set to λ = 1 for simplicity.
Under treatment selection, correlations between treatment and covariates make the treatment-effect bias depend on the design matrix and unknown nuisance parameters.With randomized assignment, the corresponding regression coefficients tend to be zero and the ridge estimate is nearly unbiased.
The problem is not bias alone, but that naive regularization can induce extreme treatment-effect bias beyond the analyst’s control, even with measured confounders and a correctly specified model.
In nonlinear and semiparametric settings, closed-form bias expressions are generally unavailable, but the RIC phenomenon can still be recreated under targeted selection.

4.2 Targeted selection

Targeted selection assigns treatment partly according to expected outcomes without treatment, creating a functional relationship between the propensity score and prognostic function. When selection is strong and aligned with prognosis, this relationship can amplify regularization-induced confounding.

Targeted selection assigns treatment partly according to the expected untreated outcome µ(x), with treatment probability increasing or decreasing as µ(x) changes.The paper notes this can occur when physicians preferentially treat patients with worse expected outcomes.
The propensity score is monotone in µ across multiple realizations indexed by ˜x, linking treatment assignment to prognosis.
For fixed auxiliary covariates ˜x, targeted selection makes the propensity function π(µ, ˜x) approximately monotone in the prognostic function µ.If the relationship is strictly monotone, π is invertible in µ for each ˜x.
In the linear formulation, RIC depends on the propensity-function coefficients, prognostic-function coefficients, and selection strength measured by Var(Z | x) = Var(ν).
The response surface can be rewritten with an arbitrary bias parameter b, showing how prior preference for transformed prognostic coefficients can shift the posterior treatment effect toward τ + b.
Bias is largest when confounding is strong and selection is targeted, particularly when propensity and prognostic coefficients point in the same direction.
The two-dimensional example uses a prognostic threshold in x1 − x2 while treatment probability increases with µ, producing targeted selection.

Example 1: d = 2, n = 250, homogeneous effects

In a simulated two-covariate setting with strong targeted selection, standard BART estimates treatment effects poorly despite a relatively favorable sample and noise regime. Allowing the model to use estimated propensity scores substantially improves estimation and interval coverage.

The example defines prognosis through a critical threshold in x1 − x2 and assigns medication increasingly to patients with higher diagnostic markers.
A modified BART prior allowing splits in an estimated propensity score performs markedly better on bias, 95% credible-interval coverage, and RMSE.The improvement primarily comes from including an estimate of π as a covariate.
Standard BART exhibits high bias, high RMSE, and poor coverage of 95% credible intervals in this setting.
Tree-based regularization struggles with the diagonal shelf because many axis-aligned splits are needed, while a treatment split can substitute for those prognostic splits under strong confounding.This mismatch makes the treatment effect prone to overstatement.
The example is deliberately stylized to be difficult for tree-based models, and other models may experience less RIC in this particular data-generating process.The broader text cautions that informative, sparse, or nonparametric priors can face similar effects.

4.3 Mitigating RIC with covariate-dependent priors

Including an estimated propensity score as a covariate in regularized response-surface models mitigates regularization-induced confounding and can improve finite-sample treatment-effect estimation.

The propensity score should be penalized like the treatment variable so regularization does not misattribute propensity-related outcome variation to treatment.
An estimated propensity score captures the observed association between treatment and covariates, making treatment conditionally independent of covariates given that score.
Adding the estimated propensity score to BART dramatically mitigates regularization-induced confounding in nonlinear models.
Incorporating the estimated propensity score into response-surface estimation improves average treatment-effect estimation in finite samples and also benefits heterogeneous conditional effects.

5 Regularization for heterogeneous treatment effects: Bayesian causal forests

The Bayesian causal forest reparameterizes the response surface into separate prognostic and treatment-effect functions, enabling direct, independent regularization of heterogeneous effects while preserving flexible modeling.

5 Regularization for heterogeneous treatment effects: Bayesian causal forests: A direct BART prior on the response surface lacks explicit control over how treatment effects vary with treatment assignment.
5 Regularization for heterogeneous treatment effects: Bayesian causal forests: Independent priors on the prognostic function µ and treatment-effect function τ allow treatment-effect regularization to be controlled separately from prognostic modeling.
5 Regularization for heterogeneous treatment effects: Bayesian causal forests: The model writes the response as a prognostic function plus a treatment-effect function, with the conditional mean contrast equal to τ(x).
5 Regularization for heterogeneous treatment effects: Bayesian causal forests: The parameterization interpolates between separate treated-control models and models treating treatment assignment as just another covariate.
5 Regularization for heterogeneous treatment effects: Bayesian causal forests: The authors expect treatment effects to be relatively small and simple, motivating stronger regularization of τ than of µ.
5.3 Data-adaptive coding of treatment assignment: Treatment coding can materially change posterior inferences on τ(x), especially when treated and control outcomes have very different marginal variances.

6 Empirical evaluations

Simulation studies and ACIC competitions evaluate BCF against tree-based, random-forest, and linear alternatives across linear, nonlinear, and targeted-selection settings. BCF performs well across scenarios, with especially favorable estimation error and shorter intervals, while coverage can slightly favor ps-BART.

Simulation studies: BCF performs well across a wide range of simulated heterogeneous-treatment-effect scenarios.The evaluations compare BCF with generalized random forests and a linear model, alongside separate ACIC challenge analyses.
Simulation studies: The simulations vary treatment-effect homogeneity, conditional-expectation linearity, and sample size, evaluating ATE and CATE estimation error, coverage, and interval length.Results use 200 independent replications for each data-generating process.
Simulation studies: BCF or ps-BART benefit dramatically from explicitly protecting against regularization-induced confounding, while BART-(f0, f1) and causal random forests perform subpar.All methods improve with a larger sample, and BCF priors are especially helpful at smaller sample sizes.
Simulation studies: BCF improvements over ps-BART are more pronounced in nonlinear settings, with usually comparable coverage and notably shorter average intervals.The linear model dominates when correctly specified but fares extremely poorly when misspecified.
ACIC 2016 competition: In ACIC 2016, BCF has the lowest estimation error for CATE and SATT among the compared methods, although differences are small relative to variation across datasets.Causal random forests perform notably worse than BART-based methods on every metric.
ACIC 2016 competition: For ACIC 2016 SATT intervals, ps-BART has 88% coverage versus BCF's 82%, while ps-BART intervals are 45% longer and vanilla BART intervals about 55% longer than BCF's.The reported nominal level is 95%.
ACIC 2017 competition: In ACIC 2017, ps-BART and BCF have the lowest estimation error for CATE and SATE, with ps-BART slightly better in coverage and BCF much shorter intervals.The 2017 datasets use smaller, less variable effects and explicitly incorporate targeted selection.

7 The effect of smoking on medical expenditures

The smoking-expenditure reanalysis finds that BCF and BART predict responses similarly but produce meaningfully different treatment-effect estimates. BCF indicates age-related heterogeneity and a lower average treatment effect than BART, despite uncertainty in most individual effects.

Data and treatment definition: Smoking is analyzed among smokers using heavy lifetime smoking, defined as more than 17 pack-years, as the treatment.The response is the natural logarithm of annual medical expenditures, so treatment effects correspond to multiplicative expenditure effects.
Model comparison: The two models have 0.98 correlation in expected responses but only 0.70 correlation in individual treatment effects.Thus, they concur on nonlinear prediction while differing substantially in estimated heterogeneous effects.
Heterogeneous effects: BCF estimates show a strong age trend and suggest that smoking has a pronounced impact on health expenditures among younger people.The propensity–prognostic score relationship is monotone, suggestive of targeted selection.
Heterogeneous effects: Although most individual treatment-effect intervals contain zero, posterior subgroup comparisons provide relatively strong evidence of heterogeneity.The posterior difference between men younger than 46 and women over 66 is almost entirely above zero, with young men having a higher estimated subgroup ATE.
Inference caveat: The analysis uses an informal posterior exploration strategy and does not further formalize its decision-theoretic extension.The authors describe the resulting summaries as valid Bayesian inferences while noting that formalization is not explored further.

8 Discussion

The discussion presents BCF as a response to regression-induced confounding and uncontrolled treatment-effect regularization, while acknowledging uncertainty about propensity-score specification and limited theory. Simulations using realistic data-generating processes support its finite-sample performance in the examined regimes.

Discussion: BCF includes an estimated propensity score in the outcome model, inducing a prior that discourages attributing confounder effects to treatment.The discussion frames this as a solution to BART’s susceptibility to regression-induced confounding.
Discussion: BCF separately models prognostic effects and treatment effects, allowing treatment-effect shrinkage to be controlled independently.Including the propensity score among splitting dimensions changes which functions are favored by the prior without materially changing its support.
Discussion: The propensity score is useful for regularization even in high-dimensional regression, because omitting it can add bias to treatment-effect estimates.The paper argues that the additional covariate can be offset by bias reduction while retaining complex response-surface modeling.
Discussion: The propensity-score adjustment depends on an accurate estimate, motivating sensitivity analysis across alternative propensity models when predictive performance does not distinguish them.Candidate propensity models can be evaluated by cross-validation for predicting treatment assignment.
Discussion: BCF differs from doubly robust estimation because its motivation is regularization that avoids regression-induced confounding, not robustness to partial misspecification or parametric convergence rates.The distinction concerns the purpose of the parameterization and propensity-score inclusion.
Discussion: The paper acknowledges nascent theory and no theory addressing frequentist coverage of Bayesian credible intervals.It nevertheless positions BCF as another supervised-learning algorithm for heterogeneous-effect estimation.
Discussion: Realistic simulation designs were intended to assess finite-sample operating characteristics rather than showcase the method’s strengths.The simulations also highlighted that methods with available asymptotic theory did not always perform as theory would suggest in finite samples.