Source-linked AI summary

Graph-Split Bayesian Causal Forest for Spatial Heterogeneous Treatment Effect Estimation

Shuren He, Huiyan Sang, Ligang Lu

arXiv:2609.36046v1stat.MEstat.ML

TL;DR

Spatial observational studies need methods that account for spatial confounding while estimating heterogeneous treatment effects, but axis-aligned causal BART splits are poorly suited to spatial structures. GSBCF combines graph-split trees with Bayesian causal forest propensity-score regression and an informed posterior sampling algorithm. Simulations and a county-level COPD study report improved estimation accuracy and uncertainty calibration over existing causal BART methods.

  • Problem

    Axis-aligned splits in causal BART are poorly suited to spatial structures, while spatially varying treatment effects and spatially structured confounding must be addressed in observational studies.

    Source: §1 · pp. 2–3
  • Method

    GSBCF integrates graph-split BART with Bayesian causal forest propensity-score regression and uses graph-guided rules to model spatial confounding in prognostic and treatment-effect functions.

    Source: §1 · p. 3
  • Results

    Simulations and a county-level COPD study show substantially improved estimation accuracy and uncertainty calibration over existing causal BART methods.

    Source: §1 · p. 3
  • Takeaways & Limitations

    Graph-guided tree partitions provide a spatial structure-aware Bayesian approach to estimating heterogeneous treatment effects in spatial observational studies.

    Source: §1 · p. 3

    Residual confounding may remain when latent confounders have weak or no spatial structure or the spatial resolution is too coarse to capture small-scale dependence.

    Source: §2 · p. 6

Abstract

from arXiv · show

In spatial observational studies, treatment assignment and outcomes often exhibit spatial dependence patterns, and treatment effects may vary across space and subpopulations due to both measured and unmeasured spatially structured confounders. Accounting for spatial dependence while estimating heterogeneous treatment effects (HTEs) is a central task in spatial causal inference. Causal Bayesian additive regression tree methods are popular nonparametric methods for modeling and estimating HTEs. Despite their flexibility and uncertainty quantification, the axis-aligned split rules often adopted in these models are not suitable for modeling spatial structures. We propose a spatial structure-aware Bayesian nonparametric method, called Graph-Split Bayesian Causal Forest (GSBCF), that integrates graph-split Bayesian additive regression trees (GS-BART) with the Bayesian causal forest propensity-score regression framework for spatial heterogeneous causal inference. Spatial confounding is accommodated through graph-guided split rules in modeling decision trees of the prognostic and HTE functions. We develop an efficient informed proposal sampling algorithm for posterior computation, enabling full Bayesian inference of the spatial conditional average treatment effect function. Simulations and a real data study demonstrate substantially improved estimation accuracy and uncertainty quantification over existing causal BART methods.

1. Introduction

Spatial causal inference must account for spatially structured confounding while estimating treatment effects that vary across locations and subpopulations. GSBCF combines graph-guided tree partitions with Bayesian causal forest propensity-score regression; the framework defines spatial causal estimands and states the assumptions under which they are identified.

1. Introduction

  • Existing spatial methods often use low-dimensional preselected covariates and assume constant effects, while axis-aligned causal BART splits are inefficient for spatial structures and irregular domains.These limitations constrain modeling of complex confounding and heterogeneous treatment effects.Source: §1 · pp. 2–3
  • GSBCF replaces axis-aligned splits with graph-guided partitions in a Bayesian causal forest, modeling spatial propensity, prognostic, and treatment-effect functions.Its graph-guided rules respect spatial contiguity and domain topology while accounting for measured covariates and unmeasured spatial factors.Source: §1 · p. 3

2. Causal inference framework for spatial data

  • The framework defines CATE as the expected treated-minus-control potential outcome conditional on covariates and location, and also defines spatially marginalized CATE and sample ATE.The supplied passages do not give the full averaging formula for the marginalized CATE or ATE.Source: §2 · p. 4
  • Identification relies on spatial latent ignorability and positivity, while the framework assumes no interference across spatial units under SUTVA.Residual confounding may remain if latent confounders lack spatial structure or the spatial resolution is too coarse; additional covariates, finer-resolution data, or spatiotemporal replication may help.Source: §2 · pp. 5–6
  • The outcome model combines a prognostic function of covariates, location, and estimated propensity score with a centered treatment term whose coefficient is the spatial CATE.Both prognostic and treatment-effect functions depend jointly on covariates and location, absorbing spatial confounder effects without explicitly modeling latent factors.Source: §2 · pp. 7–8

3. Nonparametric prior on latent spatial regression functions by GS-BART

GS-BART models spatial regression functions with graph-guided tree partitions rather than relying only on conventional axis-aligned splits. Its graph rules can preserve spatial contiguity, allow covariate–space interactions, and vary across trees.

3. Nonparametric prior on latent spatial regression functions by GS-BART

  • GS-BART represents a function as an additive ensemble of tree weak learners, with samples in each leaf sharing a constant mean.The tree prior controls splitting by node depth and selects valid graph-and-edge rules that produce nonempty children.Source: §3 · pp. 9–10
  • GS-BART replaces conventional single-covariate threshold splits with partitions guided by graph structure, addressing limitations of separable spatial models that impose stationarity or global smoothness.Separable covariate and spatial effects may also fail to capture their interactions.Source: §3 · p. 9
  • Removing an edge from a candidate arborescence splits its vertices into connected components, inducing recursive partitions of the attached samples.For numerical covariates, chain-graph edge removal corresponds to a standard threshold split; spatial graphs induce contiguous partitions respecting domain geometry.Source: §3 · p. 10
  • A tree using both spatial and covariate graph rules makes the fitted function nonseparable, enabling interactions between measured covariates and unmeasured spatial effects.Candidate graph sets can vary across trees, so the ensemble is not tied to one spatial binning and can form more expressive boundaries.Source: §3 · pp. 10–11

4. Spatial Causal Inference Method

GSBCF models prognostic and treatment-effect functions with spatially informed tree ensembles and estimates treatment effects through a Bayesian causal forest framework. Its fitting algorithm updates the forests by backfitting and uses informed proposals for tree sampling.

4.1 Bayesian Spatial Causal Forest Model

  • GSBCF represents the prognostic, treatment-effect, and treatment-assignment functions with GS-BART ensembles, including estimated propensity scores in the prognostic forest.The propensity scores are estimated using a GS-BART logistic model and included through an additional covariate chain graph.Source: §4.1 · pp. 11–12
  • Stronger depth regularization is imposed on treatment-effect trees than on prognostic trees, reflecting the assumed simpler treatment-effect heterogeneity.The specified split-prior settings are α = 0.95, β = 2 for prognostic trees and α = 0.25, β = 3 for treatment trees.Source: §4.1 · p. 12

4.2 Model Fitting Algorithm

  • Backfitting updates each tree conditional on the other tree components and residual variance, using partial residuals under a varying-coefficient likelihood.The algorithm separately updates the prognostic and treatment forests before updating the residual variance.Source: §4.2 · p. 13
  • Posterior CATE inference sums post-burn-in treatment-tree samples, while ATE estimates average CATE over the empirical covariate distribution.Posterior predictive draws at new covariate and spatial values are obtained by assigning test points to candidate-graph vertices and tree leaves.Source: §4.2 · pp. 13–14

4.3 Sampling of decision tree parameters

  • Informed Importance Tempering reweights neighboring split/merge proposals to provide rejection-free tree sampling, avoiding the low-acceptance problem of random proposals.A vectorized, recursive calculation exploits ordered arborescence structure to evaluate candidate split-move likelihood ratios efficiently.Source: §4.3 · pp. 14–16

5. Numerical Studies

The simulations compare GSBCF with Bayesian tree-based and spatial causal benchmarks on lattice and irregular U-shaped domains, under homogeneous and heterogeneous treatment effects. GSBCF performs best across the four DGPs on RMSE and CRPS, with close-to-nominal coverage for heterogeneous-effect intervals.

5.1 Simulation studies

  • The simulations include homogeneous-effect scenarios on a square lattice and a U-shaped domain, plus heterogeneous-effect scenarios on both domains.The lattice scenarios use 30 × 30 units; the irregular-domain scenarios use 800 spatial locations.Source: §5.1 · pp. 16–17
  • Benchmarks include BCF, XBCF, standard and separate-outcome BART, propensity-score BART, IPW, OPW, and a spatial GP causal model.GP-Causal is restricted to homogeneous treatment effects and cannot be applied to DGP2 or DGP4.Source: §5.1 · pp. 17–19
  • Across all four DGPs, GSBCF achieves the best RMSE and CRPS; its nominal 95% pointwise CATE intervals have coverage of 0.96 in DGP2 and 0.97 in DGP4.For heterogeneous effects, GSBCF also yields the most accurate CATE estimates, especially on the complex U-shaped domain.Source: §5.1 · pp. 19–20, Table 1 · p. 19, Table 2 · p. 20
  • On the U-shaped domain, GSBCF preserves more local CATE variation and structural detail than XBCF, although both estimates are somewhat oversmoothed.Figure 1 displays the true DGP4 spatial surface and GSBCF's estimated-surface mean absolute error.Source: §5.1 · p. 20, Figure 1 · p. 21

6. Real Data Analysis of County Level Chronic Obstructive Pulmonary Disease

The analysis estimates the effect of rural residence on 2023 age-adjusted COPD prevalence across contiguous U.S. counties, using county covariates and spatial methods. GSBCF estimates a more negative average effect than BCF and reveals spatial variation, though the Texas regional contrast remains uncertain.

6. Real Data Analysis of County Level Chronic Obstructive Pulmonary Disease

  • The study uses 2023 age-adjusted COPD prevalence and classifies Rural–Urban Continuum Codes 1–3 as non-rural and 4–9 as rural, with county-level SVI covariates.The analysis covers contiguous U.S. counties and excludes Florida and Connecticut because of data missingness and inconsistency.Source: §6 · p. 21
  • Estimated propensity scores show substantial overlap for both GSBCF and BCF except among counties easily classified as urban; Texas marginalized CATEs average predictions across county covariate profiles at each fixed location.Source: §6 · p. 22
  • GSBCF estimated an ATE of -0.54 (95% Credible Interval: [-0.66,-0.43]), compared with -0.19 (95% Credible Interval: [-0.26,-0.13]) for BCF.Both estimates indicate a positive causal effect of rural residence on COPD prevalence.Source: §6 · p. 22
  • IPW estimated an Urban-versus-Rural effect of (-0.84) (95% CI: (-2.53,0.86)), while the authors describe GSBCF’s ATE as substantially larger in magnitude.Source: §6 · p. 23
  • GSBCF estimates show a smoother, stronger spatial pattern than BCF, with smaller-magnitude rural effects around Austin, Houston, and Dallas.The authors note that greater healthcare availability in these metropolitan areas may be associated with weaker rural–urban disparities.Source: §6 · p. 23
  • In Texas, the spatially marginalized CATE tends to be less negative in western counties, but the west-minus-east contrast is uncertain: 0.05 (95% credible interval: [−0.03, 0.27]) with probability 0.66 that the western average is higher.The pattern may suggest a more pronounced rural-associated increase in eastern Texas, but the reported contrast does not establish that difference.Source: §6 · p. 23

7. Conclusion

The conclusion presents GSBCF as a graph-split Bayesian causal forest for spatially heterogeneous effects and reports an informed importance-tempering sampler for posterior exploration. The framework focuses on Gaussian outcomes, with several treatment, outcome, and interference extensions left for future work.

7. Conclusion

  • GSBCF uses graph-split decision rules in a Bayesian causal forest to estimate spatially heterogeneous treatment effects in observational studies with complex spatial structure.The framework models spatial dependence and heterogeneity in treatment effects, propensity scores, and prognostic functions.Source: §7 · pp. 23–24
  • Figure 2 depicts COPD prevalence, Texas spatially marginalized HTEs, and predictive HTE estimates for GSBCF and BCF across U.S. counties.Source: Figure 2 · p. 24, §7 · p. 24
  • The informed importance-tempering algorithm is presented as enabling efficient posterior exploration over complex tree spaces.Source: §7 · p. 24
  • The paper focuses on Gaussian outcomes; non-Gaussian outcomes, continuous or multi-group treatments, and spatial interference are identified as directions for future development.The proposed non-Gaussian extension would follow a generalized regression framework, while continuous or multi-group treatments would require generalized propensity-score regression.Source: §7 · p. 25

S1. Details of Model Fitting Algorithms

The supplement describes data-calibrated priors and a graph-split tree-fitting procedure. Candidate graph computations can be parallelized, with recursive summary-statistic calculations and informed importance tempering supporting tree proposals.

S1.1 Hyperparameter choice

  • Residual-variance prior hyperparameters follow BART-style data-dependent calibration, with shape fixed at ν = 3 and scale selected using the response’s sample standard deviation.The calibration places most prior mass below a rough data-based estimate of residual scale.Source: S1.1 Hyperparameter choice · p. 27
  • The leaf-weight variance scale adapts to the response scale and the number of trees for each latent function, while the variance itself is learned from its posterior.Source: S1.1 Hyperparameter choice · p. 27

S1.2 Parallel and recursive algorithm

S1.3 Propensity-score estimation

The propensity-score model represents treatment assignment with a GS-BART ensemble over candidate graph sets. Posterior treatment probabilities are averaged to estimate scores, which are then used as fixed predictors in the prognostic forest.

S1.3 Propensity-score estimation

  • The treatment-assignment function is represented by a GS-BART ensemble, with candidate graph sets combining covariate chain graphs and spatial adjacency-based arborescences.This permits treatment assignment to depend on observed covariates, spatial location, and their interactions.Source: S1.3 Propensity-score estimation · p. 29
  • The propensity-score estimate averages treatment probabilities over the posterior on the probability scale.Source: S1.3 Propensity-score estimation · p. 29
  • The estimated score is added to the prognostic forest but excluded from the treatment-effect forest, and the primary analysis treats it as fixed when fitting the outcome model.Source: S1.3 Propensity-score estimation · p. 29

S2. Details of candidate graph set construction

The candidate graph construction bins unstructured numerical observations into an ordered chain and spatial areal units into connected subgraphs, then links adjacent bins.

S2. Details of candidate graph set construction

  • Spatial areal units are grouped into connected subgraphs by community detection, and bins are linked when at least one original adjacency connects them.The resulting binned graph defines the spatial grouping structure and assigns training and testing samples to vertices.Source: S2. Details of candidate graph set… · p. 30
  • Unstructured numerical observations are sorted and binned into an ordered chain graph, with adjacent bins connected by directed edges.Cut points define the bins, and the number of vertices is bounded by the number of observations.Source: S2. Details of candidate graph set… · p. 30
  • Table S1 reports the numbers of arborescences and vertices for spatial or network-structured and unstructured features in each numerical example.Source: Table S1 · p. 30, S2. Details of candidate graph set… · p. 30

S3. Details of data generation process

The supplementary analyses describe convergence diagnostics and sensitivity experiments, reporting efficient sampling and robustness to propensity-score and outcome-model choices.

S3. Details of data generation process

S4.1 MCMC convergence analysis

S4.2 Sensitivity Analysis

  • Including the estimated propensity score reduces mean CATE RMSE and CRPS by approximately 16% and 20%, respectively, with further improvements using the true propensity score.The sensitivity experiment uses 50 paired replications while holding covariates, locations, and fitting settings fixed.Source: S4.2 Sensitivity Analysis · p. 32
  • Under mild SUTVA violations, GSBCF remains robust to outcome-model misspecification, performing only slightly worse than the oracle benchmark and outperforming competing causal models.The proposed explanation is that the prognostic component can capture combined spatial effects through spatial graph-guided additive trees.Source: S4.2 Sensitivity Analysis · p. 32

S4.3 Coverage of CATE credible intervals

GSBCF's credible intervals achieve near-nominal coverage in the reported DGPs, while additional simulations find strong accuracy and mixing under smooth and mixed confounding.

S4.3 Coverage of CATE credible intervals

  • For nominal 95% pointwise CATE credible intervals, GSBCF achieves coverage close to the nominal level in DGP2 and DGP4, whereas competing methods exhibit undercoverage.Empirical coverage is the proportion of true unit-level CATEs contained in nominal intervals, averaged across simulation replications.Source: S4.3 Coverage of CATE credible intervals · pp. 32–33, Table S5 · p. 33

S4.4 Performance under the Smooth and Mixed DGPs

  • GSBCF achieves the lowest RMSE and CRPS in both smooth and mixed settings, with slightly conservative CATE coverage; its advantage over BCF narrows under mixed confounding.The smooth-setting results indicate its advantages are not confined to piecewise-constant confounding surfaces.Source: S4.4 Performance under the Smooth and… · p. 34, Table S6 · p. 34
  • GSBCF has higher ESS per 100 sweeps than BCF under both settings, with comparable runtimes.Under Smooth, ESS is 42.8 versus 6.4 and runtime is 772.7 versus 676.2 seconds; under Mixed, ESS is 45.8 versus 7.2 and runtime is 776.3 versus 772.9 seconds.Source: S4.4 Performance under the Smooth and… · p. 34

S5. Additional real-data details and results

Loading 2609.36046v1…