Local Gaussian process extrapolation for BART models with applications to causal inference

We clarify notations throughout the paper. Let $\mathrm{x}=(x^{(1)},\dots,x^{(p)})$ denote a $p$-dimensional covariate vector, or regressors. Capital letter $\mathbf{X}=(\mathrm{x}^{\prime}_{1},\dots,\mathrm{x}^{\prime}_{n})^{\prime}$ denotes the $n\times p$ predictor matrix, and $\mathrm{y}=(y_{1},\dots,y_{n})$ is a vector of target values for supervised learning. First, we demonstrate our approach under the standard regression setting,

where we assume the residual term $\epsilon$ is a Gaussian noise term with mean zero and variance $\sigma^{2}$.

Let $C_{n,\alpha}(\mathrm{x})$ denote the prediction interval for a test point $\mathrm{x}$ with confidence level $1-\alpha$ given $n$ predicted values. Notations $q^{-}_{n,\alpha/2}\{y_{i}\}$ and $q^{+}_{n,\alpha/2}\{y_{i}\}$ are the quantile functions for the $\lfloor\alpha/2(n+1)\rfloor$-th smallest value and the $\lceil(1-\alpha/2)(n+1)\rceil$-th smallest value among a set of values $\{y_{1},\dots,y_{n}\}$ respectively.

The BART model (Chipman et al., 2010) assumes that the unknown function $f(\mathrm{x})$ in the regression model (1) can be approximated by a sum of regression trees, i.e.

where $\mathrm{T}_{l}$ represents a tree structure, which is a set of split rules partitioning the covariate space, and $\mathrm{\mu}_{l}$ is a vector of leaf parameters associated with the leaf nodes in tree $\mathrm{T}_{l}$.

Figure 2 illustrates a binary regression tree. The tree consists of a set of internal decision nodes that partition the covariate space to a set of terminal nodes (leaf nodes, say $\{\mathcal{A}_{1},\cdots,\mathcal{A}_{B}\}$). The tree function $g(\mathrm{x},\mathrm{T}_{l},\mathrm{\mu}_{l})$ is essentially a step function due to the constant leaf parameter. It assigns point $\mathrm{x}$ with the leaf parameter of the leaf node that it falls into in tree $\mathrm{T}_{l}$.

Trees are known to be prone to overfitting due to their high flexibility. Thus, proper regularization is necessary to achieve good out-of-the-sample performance. BART assigns a regularization prior to the tree structure that strongly favors weak or small trees. The tree prior $p(\mathrm{T}_{l})$ specifies the probability for a node to split into two child nodes at depth $d$ to be

which decreases exponentially as the tree grows deeper, implying strong regularization on the size of the tree. Chipman et al. (2010) suggest choosing $\alpha=0.95$ and $\beta=2$. The prior of each leaf parameter $p(\mu_{lb})$ is assumed to be independent normal with variance $\tau$, i.e. $\mu_{lb}\sim N(0,\tau)$. The prior of the residual variance $\sigma^{2}$ is set to be $\text{inverse-Gamma}(a,b)$.

The ensemble of trees is fitted by Bayesian backfitting and MCMC sampling schemes. Let $\mathrm{T}_{-l}$ denotes the set of all trees except $\mathrm{T}_{l}$, and similary define $\mathrm{\mu}_{-l}$. Note that the conditional posterior $p(\mathrm{T}_{l},\mathrm{\mu}_{l}\mid\mathrm{T}_{-l},\mathrm{\mu}_{-l},\sigma^{2},y)$ depends on other trees and parameters only through the residuals

Furthermore, one can integrate out the conjugate prior of $\mathrm{\mu}_{l}$ and derive the posterior of a tree $\mathrm{T}_{l}$ in closed form, i.e.,

This allows the conditional posterior $\mathrm{T}_{l}\mid\mathrm{r}_{l},\sigma^{2}$ and $\mathrm{\mu}_{l}\mid\mathrm{T}_{l},\mathrm{r}_{l},\sigma^{2}$ to be drawn separately and sequentially for $l=1,\dots,L$. After all trees are grown (we call it a sweep), the residual variance $\sigma^{2}$ is then sampled from the full conditional $\sigma^{2}\mid\mathrm{T}_{1},\cdots,\mathrm{T}_{L},\mathrm{\mu}_{1},\cdots,\mathrm{\mu}_{L},y$.

The original BART model (Chipman et al., 2010) draws trees from the posterior using a random walk Metropolis-Hastings (MH-MCMC) algorithm. Per iteration, the algorithm randomly proposes a single growing or pruning procedure to each tree and accepts or rejects it according to the MH ratios. Since the proposal in each iteration only makes a small modification, one can expect it is not efficient for the MH-MCMC algorithm to explore the posterior model space. The experiments in (Chipman et al., 2010) typically use $200$ burn-in steps and $1000$ iterations with $200$ trees. Furthermore, the algorithm does not scale well with large sizes of data observations or covariates.

XBART (He et al., 2019; He and Hahn, 2021), shortened for accelerated Bayesian additive regression trees, is a stochastic tree ensemble method inspired by BART. The essential idea of XBART is to avoid the expensive MH-MCMC computation but grow completely new trees recursively (like CART). At each node, the model considers a similar set of cutpoints as BART and uses the marginal likelihood to sample split criteria proportionally. Therefore it is much more efficient to explore the posterior of BART and scales well. Our approach is designed for both the framework of XBART and BART, while we will demonstrate the performance using XBART for computational efficiency.

To obtain the posterior prediction interval from XBART, we take the sampled trees as draws from a standard Bayesian Monte Carlo algorithm. Suppose the algorithm draws $S$ sweeps with $S_{0}$ burn-out iterations and $L$ trees per forest. The predicted posterior mean and variance at iteration $s$ are defined as $\widehat{f}_{s}(\mathrm{x})=\sum_{l=1}^{L}g_{l}(\mathrm{x},\widehat{\mathrm{T}}_{l}^{(s)},\widehat{\mu}_{l}^{(s)})$ and $\widehat{\sigma}_{s}^{2}$ respectively. The final prediction of target value $\widehat{y}_{s}$ is sampled from $N(\widehat{f}_{s}(\mathrm{x}),\widehat{\sigma}_{s}^{2})$. Furthermore, the prediction interval of target level $1-\alpha$ is given by quantiles of the predictions as follows

Note that this definition of posterior prediction interval is based on the prediction of every single tree with a constant leaf parameter, i.e., it still suffers from the extrapolation problem when the testing data is out of the range of the training.

This section reviews Gaussian process regression, a critical ingredient of the proposed XBART-GP extrapolation. Gaussian process regression is a non-parametric kernel-based Bayesian probabilistic model. It can interpolate and extrapolate as the covariance defines the model’s behavior over its full domain. Essentially, the specified covariance function characterizes the linear dependence of the response at pairs of points as a function of some measure of distance between those points in covariate space. For a textbook treatment of the Gaussian process, see Rasmussen (2003).

More specifically, Gaussian process regression assumes $\{f(\mathrm{x}),\mathrm{x}\in\mathbb{R}^{p}\}$ is a set of random variables that follows the Gaussian process with some mean function $\mu(\mathrm{x})$ and covariance function $k(\mathrm{x},\mathrm{x}^{*})$. Let $(\mathbf{X},\mathrm{y})$ denote a set of training data, and $\mathbf{X}^{*}$ denote the input of testing data. Following the standard regression setting in Equation 1, the joint distribution of function values $f(\mathbf{X}^{*})$ and $\mathrm{y}$ is given by

where $\mu(\cdot)$ is a regression function, and the covariance matrices are calculated by pre-specified kernel function $k(\mathrm{x},\mathrm{x}^{*})$

The prediction of the outcome $f(\mathbf{X}^{*})$ can be drawn from the conditional distribution $f(\mathbf{X}^{*})\sim N(\mu_{\mathbf{X}^{*}\mid\mathbf{X}},\Sigma_{\mathbf{X}^{*}\mid\mathbf{X}})$ with conditional mean and covariance:

Consider a fitted BART forest $\{\widehat{\mathrm{T}}_{l}\mid 1\leq l\leq L\}$ with $L$ trees. Let $\{\mathcal{A}_{l1},\cdots,\mathcal{A}_{lB_{l}}\}$ denote the covariate space partitioned by the $l$-th tree $\widehat{\mathrm{T}}_{l}$ with $B_{l}$ leaf nodes, and $\widehat{\sigma}^{2}$ for the estimated residual variance. For the $l$-th tree, suppose the test point $\mathrm{x}$ falls in a leaf node $b$ with covariate space $\mathcal{A}_{lb}$. Let the training data that falls in this leaf node and its residuals at the current tree (defined by Equation (4)) be noted as $\mathbf{X}^{\text{tr}}$ and $\mathrm{r}^{\text{tr}}$. Similarly, $\mathbf{X}^{\text{te}}$ and $\mathrm{r}^{\text{te}}$ denote the exterior testing data in the leaf node $b$ and its to-be-predicted residuals.

Note that the tree partitioned leaf node $\mathcal{A}_{lb}$ can extend to infinity if they are at the boundary. Specifically, the range of subset of training data $\mathbf{X}^{\text{tr}}$ falling in the $b$-th leaf node of the $l$-tree forms a $p$-dimensional hypercube $\mathcal{B}_{lb}$, which is a subset of $\mathcal{A}_{lb}$. If the testing point $\mathrm{x}$ falls within the range of training data, i.e., $\mathrm{x}\in\mathcal{B}_{lb}$, we consider it as an interior point, and its prediction follows the standard XBART model. Otherwise, if $\mathrm{x}\in\{\mathcal{A}_{lb}\setminus\mathcal{B}_{lb}\}$, it is an exterior point, and their predictions require extrapolation. Notably, this definition is local. Testing data can be exterior for some of the trees and interior for others in the XBART forest.

To extrapolate the prediction of XBART, we assume that the exterior testing data $\mathbf{X}^{\text{te}}$ and training data $\mathbf{X}^{\text{tr}}$ in the same leaf node form a Gaussian process. The joint distribution of the training and testing data in a leaf node is given by

where $\mu$ is the leaf parameter, $\mathrm{J}$ is a column vector of ones, and $\mathbf{I}$ is an identity matrix. By definition, the variance of the response $y$ is assumed to be $\sigma^{2}$. Here we assume that the residuals in each tree contribute equally to the total variance.

To reflect the relative distance among covariates, we define the covariance function of the Gaussian process as the squared exponential kernel

where $\delta_{i}$ is the range of the $i$-th variable $x_{i}$ in a leaf node, $\theta$ controls the smoothness and $\tau$ determines the scale of the function. $\mathrm{x}=(x_{1},\cdots,x_{p})$ and $\mathrm{x}^{\prime}=(x_{1}^{\prime},\cdots,x_{p}^{\prime})$ are the vector of two data points. In our experiments, we use $\theta=0.1$ and set $\tau$ to be the variance of observed responses divided by the number of trees.

The conditional distribution of the residuals $\mathrm{r}^{\text{te}}$ is given by

The prediction for the residuals of the exterior testing data associated with this leaf node is drawn from the conditional normal distribution $\mathrm{r}^{\text{te}}\sim N(\mathrm{\mu}_{\mathbf{X}^{\text{te}}\mid\mathbf{X}^{\text{tr}}},\mathbf{\Sigma}_{\mathbf{X}^{\text{te}}\mid\mathbf{X}^{\text{tr}}})$ rather than the fixed leaf parameter of the tree, which is used for predicting interior points.

It can be computationally intensive to find the corresponding training data and calculate the covariance matrix for each test data observation. Following the fast algorithm GrowFromRoot introduced in He et al. (2019), we propose an efficient method that pinpoints the exterior testing data and its neighborhood training data in the same leaf node by recursively partitioning the training and testing data at the same time. The procedures are summarized in Algorithm 1 and 2.

Certain algorithm design choices to boost performance and efficiency are worth emphasizing. First, we partition the covariates into split variables $\mathrm{v}$ and active variables $\mathrm{v}_{b}$. Split variables appear along the decision path from the root to a leaf node. Active variables $\mathrm{v}_{b}$ are a subset of split variables that satisfies the following condition: If there exists a testing point $\mathrm{x}\in\mathbf{X}^{\text{te}}$ such that $\mathrm{x}$ is outside the range of $\mathcal{B}_{lb}$ on that variable in leaf node $b$, then it is considered an active variable. Using only active variables in defining the GP covariance matrix ensures that the extrapolation depends only on those variables, which is intuitive and improves time efficiency. Second, when the training data sample size is large, we subsample training data when performing the GP extrapolation. Third, to avoid the identification of exterior points being misled by any outliers, we define the hypercube $\mathcal{B}_{lb}$ in each leaf node by the $95\%$ quantile of the training data.

This section illustrates the performance of XBART-GP and compares it to various prediction inference methods in regression modeling. We carefully construct data generating processes with covariate shifts and test the model’s predictive ability on interior and exterior data points separately.

Specifically, we compare our proposed method with the latest frequentist approach Jackknife+ (Barber et al., 2021) and its cross-validation version (CV)+. For any regression model, Jackknife+ uses the leave-one-out procedure to train the model repeatedly and obtain out-of-sample residuals to construct the predictive interval. Similarly, CV+ is a computationally efficient version of Jackknife+ that uses k-fold cross-validation to obtain the model’s out-of-sample uncertainty. We apply the two methods on XBART and Random Forest respectively, as the baseline methods.

For both XBART and XBART-GP, we use the following hyperparameters num_trees $=20$, num_iterations $=100$, ${\tt Nmin=20}$, $\tau=\text{Var}(\mathrm{y})/$num_trees (where $\text{Var}(\mathrm{y})$ is the variance of the response variable in the training set). We choose $\theta_{gp}=0.1$ and $\tau_{gp}=\tau$ to be the smoothness and scale parameters of the kernel function for XBART-GP.

When applying Jackknife+ and CV+ to XBART, we pick the same number of trees and $\tau$ value, except that the number of iterations is reduced to $100$ to reduce the running time of the two approaches. For simulations of Jackknife+ and CV+ on Random Forest, called Jackknife+ RF and CV+ RF in the rest of the paper, we use the default settings provided in the scikit-learn (Pedregosa et al., 2011) package in Python.

The real signal is drawn from four challenging functions, $f$, listed in Table 1. The response variable $y$ is generated independently from $y=f(\mathrm{x})+\epsilon$ with Gaussian noise $\epsilon$. In all cases, we generate $200$ data points with covariates $x_{j}\overset{\text{iid}}{\sim}N(0,1)$ for $j=1,\dots,d=10$ as the training set and another $200$ data points with $x_{j}\overset{\text{iid}}{\sim}N(0,1.5^{2})$ as the testing set.

Since the definition of exterior point for XBART-GP depends on the tree’s structure and varies by leaf nodes, we simplify the concept in this section to evaluate performance on in- and out-of-range data. In our simulation studies, we consider any testing data outside the training data range as exterior points. Under the above settings, approximately half of the testing data is interior, and half is exterior. The performance is then reported on interior and exterior points separately. We repeat the experiments on each data generating process $10$ times and calculate the average empirical probability of coverage and the average interval length with coverage level $1-\alpha=0.9$.

Simulation studies and time efficiency comparison on XBART-GP and its baselines were conducted in Python $3.8.10$ on a Linux machine with Intel(R) Core(TM) i7-8700K CPU @ 3.70GHz processor and 64GB RAM; eight cores were used for parallelization whenever it was applicable.

Figure 3 visualizes the simulation results on linear and single index functions. The full results of four functions are summarized in the Appendix Figure 6 and Table 3. This figure compares the root mean square error (RMSE) of point prediction, coverage rate, and interval length of various approaches. Note that the RMSE of Jackknife+ and CV+ are evaluated based on the average leave-one-out prediction and cross-validated prediction, respectively.

We notice that XBART-GP has the smallest RMSE of point prediction across all functions on both interior and exterior points. For interior data, our approach is the only one that delivers nominal coverage while the interval length is not too large. Furthermore, it also has substantially greater prediction coverage than the baseline methods on exterior points. It is worth noting that XBART-GP achieves better coverage without inflating the prediction interval too much. Indeed, the interval length of XBART-GP is just slightly larger than that of Jackknife+XBART or CV+XBART but is still smaller than that of Jackknife+RF or CV+RF.

Combining Jackknife+ and CV+ with XBART, the two methods produce competitive RMSE, coverage, and interval length compared to XBART itself on interior points. However, the prediction inference methods also fail to extrapolate as the in-sample residuals can not inform uncertainty on out-of-range data.

A motivating example of an applied extrapolation problem occurs in treatment effect estimation, in particular, when one of the necessary assumptions is violated under the potential outcome framework (Imbens and Rubin, 2015) for estimating the treatment effect.

Let $Y$ denote the response variable and $Z$ denote the binary treatment variable with $Z_{i}=1$ representing the $i$-th observation being treated or $Z_{i}=0$ as being in the control group. The potential outcomes under treatment and control are denoted as $Y_{i}(1)$ and $Y_{i}(0)$ for the $i$-th unit, respectively. The observable outcome can be expressed as $Y_{i}=Z_{i}Y_{i}(1)+(1-Z_{i})Y_{i}(0)$. The treatment effect on an individual $\mathrm{x}_{i}$ is defined as $\tau(\mathrm{x}_{i})=Y_{i}(1)-Y_{i}(0)$.

Most causal inference methods using the potential outcome framework also rely on the following assumptions:

1. 1.

   Stable Unit Treatment Value Assumption (SUTVA);

2. 2.

   Ignorability assumption: $Y_{i}(1),Y_{i}(0)\perp\!\!\!\perp Z_{i}\mid\mathbf{X}_{i}$;

3. 3.

   Positivity assumption: $0<P(Z_{i}=1\mid\mathrm{x}_{i})<1$.

The positivity assumption is also known as the overlap assumption. In this study, we differentiate the covariate space into two parts, the overlap region where the positivity assumption is satisfied, and the non-overlap region where it is not. As per Zhu et al. (2021), the violation of positivity assumptions can be categorized into two categories, structural and practical violation.

Structural violations occur when it is impossible for individuals with specific covariate values to receive treatment, leading to the treatment effect for these individuals or covariate values being irrelevant. One solution is to ”trim” the population, as described in (Ho et al., 2007; Petersen et al., 2012; Crump et al., 2009), which allows for the estimation of the treatment effect only in the overlap region (where the positivity assumption holds).

Practical violation refers to situations where a subset of covariate values is only observed in the treatment or control group due to selection bias or a limited sample size. Nevertheless, the average treatment effect (ATE) over the entire population and the conditional average treatment effect (CATE) for specific covariate values are still of interest (Yang and Ding, 2018). In these cases, the strategy is to perform model-based extrapolation, extending an estimated response surface into regions where no data was observed.

Recently, a few extrapolation methods have been developed to estimate the ATE of the entire population with poor overlap. Li et al. (2018) and Li et al. (2018) proposed using overlap weights to balance the covariates for estimating the ATE. (Nethery et al., 2019) proposed a way to define overlap and non-overlap regions based on the propensity score. They then use BART to estimate the treatment effect in the overlap region and use a spline model to extrapolate the treatment effect in the non-overlap region. Zhu et al. (2023) addressed this problem by modeling the treatment effect with the Gaussian process model, which by its nature can extrapolate based on the covariate kernel.

In this section, we explore extrapolating CATE in the region of non-overlap by applying the local GP extrapolation method on the Accelerated Bayesian Causal Forest (XBCF) model (Krantsevich et al., 2022). First, we briefly review the structure of the XBCF model. Next, we formally define the overlap and non-overlap area in the ensemble trees framework. Then we describe the modifications we have made to integrate the local GP algorithm for extrapolating the treatment effect estimated by XBCF.

XBCF is a recent approach for estimating heterogeneous treatment effect based on Bayesian Causal Forest (Hahn et al., 2020). The model expresses the observable outcomes as

where $\widehat{\pi}_{i}$ is the estimated propensity score. Function $\mu(\mathrm{x}_{i},\widehat{\pi}_{i})$ aims to capture the prognostic effect and function $\tau(\mathrm{x}_{i})$ captures the treatment effect. Both functions are modeled by a sum of XBART trees, respectively.

Under the derived split criterion of XBCF, the treatment trees in $\tau(\mathrm{x}_{i})$ stop splitting when the leaf nodes only consist of either the treatment group or the control group. However, the prognostic forest $\mu(\mathrm{x}_{i},\widehat{\pi}_{i})$ is not affected by the positivity violation as it does not depend on the treatment variable $z$. In this case, both estimators will be biased when the positivity assumption is violated. However, the tree structure effectively captures the heterogeneity in the overlap region, making it ideal for extrapolating causal effects in the non-overlap region using the Gaussian process model in the leaf nodes.

The model relies on the three basic assumption in the potential outcome framework describe in the previous section. However, when the positivity assumption is violated in practice, we can not make firm conclusions about ATE or CATE in the non-overlap region. To mitigate this limitation, we make the assumption that there is no abrupt change in the treatment effect function in both overlap and non-overlap regions, which allows us to make inferences about CATE using the extrapolated predictive interval obtained from local Gaussian processes.

We apply the local GP technique on the XBCF model to extrapolate the treatment effect in non-overlap regions as follows. Given a regression tree, consider a leaf node $b$ with local covariate subspace $\mathcal{A}_{\mathbf{X}}$, let matrix $\mathbf{X}^{\text{tr}}$ denote the training data that falls into the leaf node, which can be separated into a treatment group $\mathbf{X}^{\text{tr}}_{T}$ and a control group $\mathbf{X}^{\text{tr}}_{C}$ based on the corresponding treatment variable $z$. The range of covariates in the treated and the control form two hypercubes, here we denote as $\mathcal{B}_{T}$ and $\mathcal{B}_{C}$ respectively. The overlap region in the leaf node is the intersection of $\mathcal{B}_{T}$ and $\mathcal{B}_{C}$, or $\mathcal{B}_{O}=\mathcal{B}_{T}\cap\mathcal{B}_{C}$. Consequently, the non-overlap area is $\mathcal{B}_{NO}=\mathcal{A}_{\mathbf{X}}\setminus\mathcal{B}_{O}$.

It’s important to note that in each treatment tree, the residuals in the non-overlap region are biased in direction relative to the true treatment effect. This is because of the treatment trees’ no-splitting behavior. As a result, if testing data is included in the Gaussian process, it will be extrapolated in the opposite direction. Therefore, different from XBART-GP, we assume that the testing data in the non-overlap area $\mathbf{X}^{\text{te}}_{\text{NO}}$ share a joint distribution with only the overlap training data $\mathbf{X}^{\text{tr}}_{\text{O}}$, namely

Here $\mathrm{r}^{\text{tr}}_{\text{O}}$ is a vector of partial residuals for the current treatment tree, i.e., the difference between observed response $y$ and the predictions of all other prognostic and treatment trees. $\mathrm{r}^{\text{te}}_{\text{NO}}$ is a vector of residuals we wish to predict for the non-overlap testing data. $\widehat{\boldsymbol{\sigma}}^{2}_{\mathrm{z}}$ is a diagonal matrix with diagonal elements equal to estimated variance $\widehat{\sigma}_{0}^{2}$ or $\widehat{\sigma}_{1}^{2}$ depending on the treatment variable of the training data and $L$ is the total number of prognostic and treatment trees.

The prediction algorithm for the treatment trees in XBCF-GP follows the same procedures as in Algorithm 2 except for the difference we describe above. Specifically, we replace the out-of-range hypercube $\mathcal{B}_{lb}$ in line $4$ by the overlap area $\mathcal{B}_{O}$. In line $6$, the subset of training data $\mathbf{X}^{\text{tr}}_{b}$ is sampled from the training data $\mathbf{X}^{\text{tr}}_{\text{O}}$ from the overlap area instead of the entire training set. Then in line $7$, we choose to extrapolate the testing data $\mathbf{X}^{\text{te}}_{\text{NO}}$ in the non-overlap area for XBCF-GP. Similar to the algorithm design choices made in XBART-GP, the overlap area is defined by the $95\%$ quantiles of the treated and control data for robustness concerns.

In addition to the modification we make to the XBART-GP algorithms, we also add a tweak to the original XBCF model to ensure the extrapolation works smoothly. To ensure successful extrapolation, the GP requires a minimum amount of overlap data in the leaf nodes for providing accurate estimates. While it is uncommon for the treatment trees in XBCF to split in the non-overlap area, it could still occur in rare cases and result in suboptimal extrapolation. To address this issue, we implemented a strong prior on the treatment trees to discourage splits that do not have sufficient treatment and control data in their children nodes. The amount of overlap data required in a leaf node can be controlled through the hyperparameter Nmin. This hyperparameter is critical as it determines the balance between the quality of the method’s estimates near the overlap area boundary and its extrapolation performance in the non-overlap region. In our experiments using $500$ data points, we set Nmin to $20$.

The XBCF-GP model combines the benefits of the Bayesian Causal Forest and Gaussian Process. XBCF excels at accurately and efficiently estimating homogeneous and heterogeneous treatment effects on large datasets with many covariates, provided the positivity assumption holds. The GP component, applied on a per-leaf node basis, enables extrapolation of treatment effects in regions where the positivity assumption is violated by leveraging the most relevant covariates in the surrounding area.

Our method provides a more effective way to assess the uncertainty of the estimated treatment effect in cases where the positivity assumption is violated. By using the local Gaussian process, our uncertainty estimation takes into account the distance between the testing data and the overlap region. This means that if the testing data is significantly distant from the overlap region, XBCF-GP will reflect this uncertainty through wider prediction intervals.

This section compares prediction interval coverage and length under the setting of causal inference when the positivity assumption is violated. To demonstrate the performance of XBCF-GP, we compare it with a set of alternative extrapolation approaches proposed by Nethery et al. (2019). The baseline methods we use are summarized as follows:

• •

  SBART Also known as BART-Stratified. Fit separate BART models for the treated and the control group, then estimate the treatment effect as the difference between the expected value of the treated and the control models.

• •

  UBART Untrimmed BART is implemented by Nethery et al. (2019), in which a single BART model is fitted with covariates, treatment variable, and estimated propensity score. The treatment effect is estimated by predicting the potential outcomes with the posterior predictive distributions.

• •

  BART+SPL Nethery et al. (2019) proposed this method to estimate the overlap area with BART and the non-overlap area with a spline model. The region of overlap is defined by propensity score with recommended parameters $a=0.1$ and $b=7$.

Figure 5 illustrates the performance of XBCF-GP and other baseline methods on a one-dimensional toy dataset. The independent variable $x$ is uniformly drawn from the range $[-10,10]$. The prognostic function is sine, and the treatment effect is $0.25x$. The probability for a sample being treated is $\pi(x)=\max\{0,\min\{1,0.08x+0.5\}\}$. Let $f(x)=\sin(x)+(0.25x)z$ denotes the true function with treatment variable $z\sim\text{Bern}(\pi(x))$. The response values $y$ are generated as $y=f(x)+\epsilon$, where $\epsilon\sim N(0,0.2\times\text{sd}(f))$ with the standard deviation $\text{sd}(f)$ taken over the dataset. The fitted models in the example use the same parameter setup as in the simulation study, which will be specified shortly.

The figure demonstrates the performance of different models in estimating CATE. The first three methods (XBCF, XBCF-GP, and SBART) are relatively robust to noise. However, the last two methods (UBART and BART+SPL), which directly estimate CATE using the response variable, introduce significant noise in the estimates. All three BART-based methods are susceptible to the influence of the prognostic structure, especially in the non-overlap area. XBCF dodges the bullet by setting up a separate BART forest to control for the prognostic effect, but it cannot extrapolate properly in positivity violation. XBCF-GP inherits the well-regularized structure from XBCF and extrapolates the treatment effect estimations nicely and smoothly.

For a more thorough comparison among these methods, we adopt the same data generating process used in Hahn et al. (2020) and Krantsevich et al. (2022) but modify the propensity function to create the positivity violation scenario. The covariate vector $\mathrm{x}$ consists of five variables. The first three (denoted as $x_{1},x_{2},x_{3}$) are generated from the standard normal distribution. $x_{4}$ is a binary variable, and $x_{5}$ is an unordered categorical variable with levels denoted as $1,2,3$. The prognostic function can be either linear or nonlinear:

where $g(1)=2$, $g(2)=-1$ and $g(3)=-4$. The treatment effect can be either homogeneous or heterogeneous:

In the simulation studies, we adjust the propensity function such that roughly $20\%$ of the data has a propensity score of $1$, $20\%$ of the data has a propensity score of $0$, and the rest of the data has a propensity score that is between $0$ and $1$. The propensity function $\pi(\mathrm{x})$ is given by

where $s$ is the standard deviation of $\mu(\mathrm{x})$ taken over the observed samples and $u_{i}\sim\text{Uniform}(0,1)$. We pick the value $c=0$ for the linear prognostic function and $c=3$ for the non-linear one. For any propensity score greater than $1$ or less than $0$, we set it to be $1$ or $0$, respectively. In addition, we add a Gaussian noise $\epsilon\sim N(0,0.5\times\text{sd}(\mu(\mathrm{x})+\tau(\mathrm{x})\mathrm{z}))$ with the standard deviation taken over the samples to generate the response values.

To gauge the performance of XBCF-GP across overlap and non-overlap regions, we evaluate the average treatment effect (ATE) and conditional average treatment effect (CATE) on three basic metrics: average root mean square error, coverage, and average interval length. In addition, we also compare the time efficiency of the proposed method to baseline methods. We average the results from $20$ independent replications for each scenario with $n=500$ samples.

For methods that rely on the propensity score, we use the propensity score estimated by a XBART Multinomial model with $2$ classes and $100$ iterations instead of the true propensity scores. For XBCF-GP, we set hyper-parameter $\theta=0.1$ and $\tau=\widehat{\sigma}^{2}/L_{\tau}$, where $\widehat{\sigma}^{2}$ is the estimated variance and $L_{\tau}=20$ is the default number of treatment trees in the XBCF model. We choose a different scale parameter $\tau_{gp}$ from XBART-GP as the variance of the treatment effect tends to be smaller than the total variance of the observed data in most cases. As for the other baseline methods, including XBCF, we use their default parameters for all scenarios.

The experiments were conducted in R 4.1.3 on the same machine described in the previous sections.

Table 2 summarizes the simulation results of various methods. We focus our discussion here on CATE estimation. Both XBCF methods provide the best CATE estimation in all cases and deliver competitive CATE coverage with tighter intervals than other methods. XBCF-GP can improve both the accuracy and coverage of ATE and CATE estimation based on its base model. Furthermore, we notice that the efficiency of XBCF-GP, inherited from XBART, is substantially faster than other BART-based methods.