Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.
69,261 characters · 17 sections · 44 citation commands
Local Gaussian process extrapolation for BART models with applications to causal inference
\def\spacingset#1{ {#1}} \spacingset{1}
\if00 \fi
\if10 {
} \fi
{\it Keywords:} Tree, Extrapolation, Gaussian process, Predictive interval, XBART, XBCF
\spacingset{1.5}
Tree-based supervised learning algorithms, such as the Classification and Regression Tree (CART) breiman1984classification, Random Forests breiman2001random, and XGBoost chen2016xgboost are popular in practice due to their ability to learn complex nonlinear functions efficiently. Bayesian Additive Regression Trees (BART, chipman2010bart) is the most popular model-based regression tree method; it has been demonstrated empirically to provide accurate out-of-sample prediction (without covariate shift), and its Bayesian uncertainty intervals often out-perform alternatives in terms of frequentist coverage (see chipman2010bart, Kapelner2013PredictionWM). XBART he2021stochastic is a stochastic tree ensemble method that can be used to approximate BART models in a fraction of the run-time. Throughout the paper, we will refer to BART models but will use the XBART fitting algorithm.
While tree-based methods frequently provide accurate out-of-sample predictions, their ability to extrapolate is fundamentally limited by their intrinsic, piecewise constant structure. Trees partition the covariate space into disjoint rectangular regions (leaf nodes) based on the observed covariates, then estimate a simple model (typically a constant) based on training data falling in the corresponding region. Consequently, to predict testing data outside the observed range of covariates in training, a naive tree regression model would extrapolate the constant in whatever leaf node that point falls in. Figure (ref) illustrates the potential problem of extrapolation for the traditional tree models with constant leaf parameters in one-dimensional covariate space. If the testing data lie out of the range of the training, tree models have to extrapolate using the constant leaf parameter in the corresponding leaf node regardless of how far away the testing observations are from the nearest training point.
In addition to poor point predictions, a flawed extrapolation model impairs posterior prediction intervals. Ideally, a BART model would predict essentially according to the prior at points far from the observed data; standard implementations instead never sample cutpoints beyond the training range and consequently dramatically understate the uncertainty in the unknown function in extrapolation regions. As a result, the coverage of predictive intervals from standard tree-based methods can be poor. While regression tree methods can be used in conjunction with resampling approaches to construct improved prediction intervals efron1981Jackknife, papadopoulos2008inductive, vovk2015cross,barber2021predictive, such a strategy does not solve the intrinsically poor extrapolation of regression trees.
This paper presents a novel approach to help tree methods extrapolate the exterior points and construct prediction intervals. The new method (XBART-GP) fits a Gaussian process (GP) at each leaf node of a BART tree and predicts exterior test points in the leaf by the associated GP model. Our approach differs from a treed Gaussian process gramacy2008bayesian in that the local GP extrapolation does not impact the initial BART model fit. Intuitively, model of gramacy2008bayesian is a “treed" GP, while our approach is a “GPed" tree. The local GP extrapolation defines a novel posterior predictive distribution in terms of BART posterior samples and a GP model; it is not a new likelihood for inferring the trees themselves. Instead, the BART-inferred trees are used as inputs to the novel posterior predictive specification (cf. section (ref)). Our proposed method shares a similar structure with the targeted smoothing BART (tsBART) model Starling2018tsBART, Starling2021heterogneous, which fits a Gaussian process on a single targeted covariate in each tree leaf to introduce smoothness. In contrast, our approach fits the Gaussian process on a selection of covariates to extrapolate prediction.
In simulation studies, we compare our approach with recently proposed alternatives such as Jackknife+ barber2021predictive. Our results illustrate that XBART-GP achieves nearly the desired coverage on out-of-range data while not substantially inflating the interval length.
Additionally, our new extrapolation method is illustrated with an application to causal inference. When only treated (untreated) units are observable in a particular region of covariate space, inferring the untreated (treated) potential outcomes requires extrapolating the untreated (treated) response surface into that region. Such a scenario is said to violate the positivity condition necessary for identifying causal effects from observational data. Previous work on positivity violation includes DAMOUR2021644, who focus on high-dimensional covariate spaces; nethery2019estimating, who propose a spline model to extrapolate in non-overlap regions; and zhu2023addressing, who construct a Gaussian process for the entire covariate space. In this paper, we apply the proposed local GP extrapolation approach to the Bayesian causal forest model hahn2020bayesian, nikolay2022.
The rest of the paper is organized as follows. Section (ref) reviews Gaussian processes, BART, Jackknife+, and Bayesian causal forests. Section (ref) develops the new approach in detail and presents simulation studies demonstrating its advantages over the current state-of-the-art. Section (ref) adapts the method to the causal inference context and presents simulation results showing that local GP extrapolation of BART models can handle substantial violations of the overlap condition.
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}'_1, \dots, \mathrm{x}'_n)'$ 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 chipman2010bart assumes that the unknown function $f(\mathrm{x})$ in the regression model (ref) 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 (ref) 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. chipman2010bart 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 chipman2010bart 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 chipman2010bart 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 he2019xbart, he2021stochastic, 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 rasmussen2003gaussian.
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 (ref), 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:
Traditional BART prediction intervals in Equation ((ref)) provide good coverage with tight intervals on many data generating processes he2021stochastic with exchangeable train and test data. However, predictions are less reliable when the train and test sets differ substantially in terms of their support. The intuition behind our method is to use the tried-and-true BART predictions (and intervals) for prediction points within the range of the training data but to incorporate Gaussian process extrapolation for points well outside this support. Therefore, the first step in describing the new method is to formally define the concepts of exterior points in the context of regression trees. Specifically, we will define extrapolation points locally in terms of the regression trees: a test point $\mathrm{x}'$ is an exterior point if and only if it is outside the range of the training data in the leaf node it falls in. The basic strategy of the new method is to use the leaf-specific training data to extrapolate such exterior points using a local Gaussian process. Combining these local GP predictions across trees (and across posterior samples) constitutes the main technical content of the new method, details of which are below.
Before getting into those details, a big-picture explanation may be helpful. Typically, a Bayesian posterior predictive distribution is $$f(\tilde{y} \mid y_{1:n}, \mathbf{X}_{1:n}) = \int f(\tilde{y} \mid \theta) \pi(\theta \mid y_{1:n}, \mathbf{X}_{1:n}) d\theta ,$$ which conveys the idea that future and past data are conditionally independent given the model parameters. Here, we explicitly deviate from this formulation, but only for those future points that are exterior (defined as a function of the model parameters). For such a point, our posterior predictive is defined as
$$f(\tilde{y} \mid y_{1:n}, \mathbf{X}_{1:n}) = \int f(\tilde{y} \mid \theta, y_{1:n}, \mathbf{X}_{1:n}) \pi(\theta \mid y_{1:n}, \mathbf{X}_{1:n}) d\theta $$ where the predictive distribution explicitly involves {\em both} the training data and model parameters (trees and leaf means in the case of BART). In this sense, the proposed approach does “use the data twice” but not in the construction of the posterior. Rather, our procedure is a way of combining the orthodox BART posterior with an intentionally and explicitly distinct posterior predictive model for extrapolation points. Although this is an uncommon approach that violates so-called “diachronic coherence” skyrms2006diachronic, it is in no sense illicit in that it merely amounts to specifying a user-defined predictive distribution that takes BART posterior samples as inputs. From this perspective, describing the method amounts to providing a definition of the predictive kernel, $f(\tilde{y} \mid \theta, y_{1:n}, \mathbf{X}_{1:n})$. Specifically, this will be a Gaussian process with a covariance kernel defined in terms of the trees of the BART model.
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 (ref)) 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}' = (x_1',\cdots,x_p')$ 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 he2019xbart, 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 (ref) and (ref).
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+ barber2021predictive 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 {\tt num_trees} $= 20$, {\tt num_iterations} $=100$, ${\tt Nmin = 20}$, $\tau = \text{Var}(\mathrm{y})/${\tt 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 {\tt scikit-learn} scikit-learn package in Python.
The real signal is drawn from four challenging functions, $f$, listed in Table (ref). 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 (ref) visualizes the simulation results on linear and single index functions. The full results of four functions are summarized in the Appendix Figure (ref) and Table (ref). 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.
We then evaluate the performance and time efficiency of XBART-GP and other baselines for increasing sample size.
The training and testing covariates are generated from the same distributions as in the previous section, with testing covariates having a slightly larger variance. We generate the response variable $Y$ with the linear function in Table (ref) and the standard normal error term. We estimate the average coverage, interval width, and computational time of all $6$ methods on with sample size $n = n_t = 50, 100, 150, 200, 300, 500$ over $10$ independent trials for each sample size. The target coverage level is $1 - \alpha = 0.9$.
Figure (ref) illustrates the performance comparison of XBART-GP and other baseline methods on exterior points with increasing sample size. The performance on interior points is similar and is presented in Appendix Figure (ref). Note that the proportion of exterior points generated varies according to the training size. Empirically, the percentage of exterior points decreases from $80\%$ to $40\%$ when the sample size increases from $50$ to $500$ on $10$-dimensional covariates. The x-axises in Figure (ref) only reflect the sizes of the training sets. The running time is reported as evaluating all testing data, including interior and exterior points.
On exterior points, XBART-GP has the smallest RMSE of point prediction and a stable coverage close to $90\%$ as the sample size grows. The coverage of other baseline methods drops significantly on a larger sample size. The interval of XBART-GP is slightly wider than XBART but still outperforms random forest approaches.
XBART and CV+ approaches are the most efficient in terms of running time. The run-time of XBART-GP increases almost linearly with the amount of training and testing sample size, but it is worth the time when the distribution of the test set shifts from the training set.
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 imbens2015causal 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:
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 Zhu2021PositivityViolation, 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_imai_king_stuart_2007, petersen2012, crump2009, 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_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. fanli2018a and fanli2018b proposed using overlap weights to balance the covariates for estimating the ATE. nethery2019estimating 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. zhu2023addressing 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 nikolay2022. 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 hahn2020bayesian. 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 $\ref{alg:PFR}$ 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 {\tt 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 {\tt 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 nethery2019estimating. The baseline methods we use are summarized as follows:
Figure (ref) 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 hahn2020bayesian and nikolay2022 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: \[ \mu(\mathrm{x}) =
\] where $g(1) = 2$, $g(2) = -1$ and $g(3) = -4$. The treatment effect can be either homogeneous or heterogeneous: \[ \tau(\mathrm{x}) =
\]
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 (ref) 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.
The Gaussian process extrapolation technique developed here has several advantages over standard BART extrapolation, as implemented in the most widely available software. First, whereas BART extrapolates a constant function beyond the range of the training data by fusing a Gaussian process at each leaf, “local” Gaussian process extrapolation is possible. This local GP approach automatically incorporates variable relevance, using only variables that are used as splitting rules in a given tree. Because extrapolation occurs at each tree, our procedure also generally improves prediction: GP prediction at {\em local} exterior points imposes additional smoothness even at points within the training data's convex hull.
Second, BART's prediction intervals are dominated by the residual variance parameter value, which implies that the interval width is nearly constant at any exterior point. By contrast, the local GP extrapolation method implies interval widths that expand quite rapidly the further away a point is from the training data, a desirable property that is directly inherited from the Gaussian process.
Third, the local GP approach requires no additional training time compared to a regular BART model fit. Instead, the extrapolation component is a direct modification of the posterior predictive distribution defined in terms of BART posterior samples and a specified covariance function. Broadly, our approach grew out of attempts to answer the question: how would one want to use a known BART model that fits well at interpolation points to define an extrapolation model? Our proposed solution is to graft a Gaussian process to each leaf node for the purpose of extrapolation only; posterior uncertainty in the BART model parameters is then propagated by averaging over posterior samples. The general philosophy behind this approach is that there is no need to change the likelihood model if BART adequately fits the training data, but conversely, there is no need to stick to the pure BART model for points far from the training data. Extrapolation always requires a leap of faith, and there is no guarantee that any extrapolation method will succeed, but we argue that there are conceptual and computational advantages to disentangling the extrapolation model from the interpolation model.
Finally, we demonstrate how to apply local GP extrapolation of BART models to the problem of imperfect overlap in treatment effect estimation, a common practical challenge in applied causal inference. By applying the local GP extrapolation technique to the Bayesian causal forest model, violations of the positivity assumption in practice can be handled organically. Specifically, we are able to extrapolate the treatment effect surface directly rather than composing it as the difference between two separate extrapolations (one for each treatment arm). Our simulation studies demonstrate that XBCF-GP produces treatment effect estimates with good coverage on a variety of data-generating processes. These tools will help other researchers to draw more reliable inferences from applied data which frequently exhibit imperfect overlap.
\paragraph{Acknowledgement} Jingyu He gratefully acknowledges funding for this project fully supported by the Research Grants Council of the Hong Kong Special Administrative Region, China (Project No. CityU 21504921).