Bayesian additive regression trees and the General BART model

Summary

Bayesian additive regression trees (BART) is a flexible prediction model/machine learning approach that has gained widespread popularity in recent years. As BART becomes more mainstream, there is an increased need for a paper that walks readers through the details of BART, from what it is to why it works. This tutorial is aimed at providing such a resource. In addition to explaining the different components of BART using simple examples, we also discuss a framework, the General BART model, that unifies some of the recent BART extensions, including semiparametric models, correlated outcomes, statistical matching problems in surveys, and models with weaker distributional assumptions. By showing how these models fit into a single framework, we hope to demonstrate a simple way of applying BART to research problems that go beyond the original independent continuous or binary outcomes framework.

1 ∣. INTRODUCTION

Bayesian additive regression trees (BART; Chipman et al.¹) has gained popularity in the recent years among the research community with numerous applications including biomarker discovery in proteomic studies², estimating indoor radon concentrations³, estimation of causal effects^4,5, genomic studies⁶, hospital performance evaluation⁷, prediction of credit risk⁸, predicting power outages during hurricane events⁹, prediction of trip duration in transportation¹⁰, and somatic prediction in tumor experiments¹¹. BART has also been extended to survival outcomes^12,13, multinomial outcomes^14,15, and semi-continuous outcomes¹⁶. In the causal inference literature, notable papers that promote the use of BART include Hill⁵ and Green and Kern¹⁷. BART has also been consistently among the best performing methods in the Atlantic causal inference data analysis challenge^18,19,20. In addition, BART has been making inroads in the missing data literature. For the imputation of missing covariates, Xu et al. ²¹ proposed a way to utilize BART for the sequential imputation of missing covariates, while Kapelner and Bleich²² proposed to treat missingness in covariates as a category and set up the splitting criteria so that the eventual likelihood in the Metropolis-Hasting (MH) step of BART is maximized. For the imputation of missing outcomes, Tan et al.²³ examined how BART can improve the robustness of existing doubly robust methods in situations where it is likely that both the mean and propensity models could be misspecified. Other more recent attempts to utilize or extend BART include applying BART to quantile regression²⁴, extending BART to count responses²⁵, using BART in functional data²⁶, applying BART to recurrent events²⁷, identifying subgroups using BART^28,29,30, using BART as a robust model to impute missing principal strata to account for selection bias due to death³¹, decision making and uncertainty quantification for individualized treatment regimes³², as well as competing risks³³.

The widespread use of BART has resulted in many researchers starting to use BART as a reference model for comparison when proposing new statistical or prediction methods. A few recent examples include Liang et al.³⁴, Nalenz and Villani³⁵, and Lu et al.³⁶. This growing interest for BART raises a need for an in-depth tutorial paper to help researchers better understand the method. The first portion of this paper is aimed at addressing this.

The second portion of our work revolves around extensions of BART beyond the original independent continuous or binary outcomes setup. Recent papers involved extensions of BART to semiparametric models³⁷, correlated outcomes³⁸, statistical matching problems in surveys³⁹, and more flexible outcome distributions⁴⁰. Although these papers were written separately, they share a common feature in their framework. In brief, when estimating the posterior distribution, they subtract a latent variable from the outcome and then model this residual as BART. This idea, although simple, is powerful because this can allow researchers to easily extend BART to problems that they may face in their data set without having to rewrite or re-derive the Monte Carlo Markov Chain (MCMC) procedure for drawing the regression trees in BART. We summarize this idea in a framework unifying these models that we call, the General BART model. We believe that by presenting our General BART model framework and linking it with the models in these four papers as examples, it will be a valuable tool for researchers who are trying to incorporate and extend BART to solve their research problems.

Our in-depth review of BART in Section 2 focuses on three commonly asked questions regarding BART: What gives BART flexibility? Why is it called a sum of regression trees? What are the mechanics of the BART algorithm? In Section 3, we demonstrate the superior performance of BART compared to the Bayesian linear regression (BLR) when data are generated from a complicated model. We then describe the application of BART to two real-life data sets, one with continuous outcomes and the other with binary outcomes. Section 4 lays out the framework for our General BART model that allows BART to be extended to semiparametric models, correlated outcomes, survey, and situations where a more robust assumption for the error term is needed. In each of these examples, we describe how the prior distributions are set and how the posterior distribution is obtained. We conclude with a discussion in Section 5.

2 ∣. BAYESIAN ADDITIVE REGRESSION TREES

We begin our discussion with the independent continuous outcomes BART. We argue that BART is flexible because it is able to handle non-linear main effects and multi-way interactions without much input from researchers. To demonstrate how BART handles these model features, we explain using a visual example of a regression tree. We then illustrate the concept of a sum of regression trees using a simple example with two regression trees. We next show how a sum of regression trees accommodate non-linear main and multi-way interaction effects. We provide two perspectives to show how BART determines these non-linear effects automatically. First, we discuss the BART mechanism using a visual and detailed breakdown of the BART algorithm at work with a simple example, providing the intuition for each step along the way. Second, we provide a more rigorous explanation of the BART MCMC algorithm by discussing the prior distribution used for BART and how the posterior distribution can be calculated. Finally, we show how BART handles independent binary outcomes.

2.1 ∣. Continuous outcomes

2.1.1 ∣. Single regression tree

To understand BART we first introduce regression trees. Suppose we have covariates x = (x₁,…, x₅) and outcome y. A regression tree represents the conditional expectation of y given x. We provide a visual representation of a hypothetical, simple regression tree in Figure 1. Each place where there is a binary decision split is called a node. At the top node (root), there is a condition x₂ < 100. If x₂ < 100 is true, we follow the path to the left, otherwise we follow the path to the right. Assuming that x₂ < 100 is true, we see that we arrive at a node which is not split upon. This is called a terminal node and the parameter μ₁ = 1.19 is the assigned value of E[y∣x] for any x where x₂ < 100. Suppose instead that x₂ < 100 is not true. Then, moving along the right side, another internal node with condition x₄ < 200 is encountered. This condition would be checked and, if this condition is true (false), we follow the path to the left (right). This process continues until we reach a terminal node and the parameter μ_i in that terminal node is assigned as the value of E[y∣x], where μ_i is the mean parameter of the i^th node for the regression tree. So, for example, a subject k with x_k1 = 30, x_k2 = 120, x_k3 = 115, x_k4 = 191, and x_k5 = 56 would be assigned a value of μ₂ = 2.37 for E[y∣x]. The value would be exactly the same for another subject k′ who instead had covariates x_k′1 = 130, x_k′2 = 135, x_k′3 = 92, x_k′4 = 183, x_k′5 = 10.

FIGURE 1.

Example of a regression tree g(x; T, M) where μ_i is the mean parameter of the i^th node for the regression tree.

We denote by T the binary tree structure itself - the various binary split decision rules of the form {x_q < c} versus {x_q > c}, $c\in \mathbb{R}$. The vector of parameters associated with T, i.e. the collection of terminal node parameters, is denoted by M ={μ₁,…,μ_b}, where b is the number of terminal nodes. Defining a regression tree as g(x; T, M), we can view g(x; T, M) as a function that assigns the conditional mean E[y∣x] to the parameter μ_i ∈ M i.e. μ_i = g(x; T, M) ↦ E(y∣x) via binary decision rules denoted as T. Note that we have not discussed how these binary decision rules in a regression tree are created by BART and how uncertainty about which covariate to split on and which value to split on, is quantified. We will address that when we introduce the BART priors and algorithms.

Regression tree as an analysis of variance (ANOVA) model

Another way to think of the regression tree in Figure 1 is to view it as the following analysis of variance (ANOVA) model:

$$\begin{gathered} y={\mu }_{1}I\{x_2<100\}+{\mu }_{2}I\{x_2\ge 100\}I\{x_4<200\}I\{x_3<150\} \\ +{\mu }_{3}I\{x_2\ge 100\}I\{x_4<200\}I\{x_3\ge 150\}I\{x_5<50\} \\ +{\mu }_{4}I\{x_2\ge 100\}I\{x_4<200\}I\{x_3\ge 150\}I\{x_5\ge 50\} \\ +{\mu }_{5}I\{x_2\ge 100\}I\{x_4\ge 200\}+\epsilon , \end{gathered}$$

where I{.} is the indicator function and ε ~ N(0, σ²). We can see that the term μ₁ I{x₂ < 100} corresponds to the terminal node on the top left corner of Figure 1, μ₂I{x₂ ≥ 100}I{x₄ < 200}I{x₃ < 150} corresponds to the terminal node just below μ₁ = 1.19, and so on. We can think of μ₁I{x₂ < 100} as a main effect, because it only involves the second variable x₂, while μ₂I{x₂ ≥ 100}I{x₄ < 200}I{x₃ < 150} is a three way interaction effect involving the second (x₂), fourth (x₄), and third variable (x₃). By viewing a regression tree as an ANOVA model, we can easily see why a regression tree and hence, BART, which is made up of a sum of regression trees, is able to handle main and multi-way interaction effects.

2.1.2 ∣. Formal definition

We now formally define BART. Suppose we have a continuous outcome y and p covariates x = (x₁,…, x_p). The goal is a model that can capture complex relationships between x and y, with the aim of using it for prediction. BART attempts to estimate f(x) from models of the form y = f(x)+ε, where, for now, ε ~ N(0, σ²). To estimate f(x), a sum of m regression trees is used i.e. $f(x)={\sum }_{j=1}^{m}g(x;T_j,M_j)$. Thus, BART is often presented as

$$y=f(x)+\epsilon =\sum _{j=1}^{\mathit{m}}g(x;T_j,M_j)+\epsilon .$$ (1)

where T_j is the j^th binary tree structure (recall T from Section 2.1.1) and M_j = {μ_j1,…, μ_jbj} is the vector of terminal node parameters associated with T_j. The number of trees m, is usually fixed at a large number, e.g., 50, 100, or 200. For example Bleich et al.⁴¹ suggests that 50 is often adequate. Note that some authors may also refer to the BART prior in the model setup as f ~BART. In addition, some readers may recognize that using y = μ + f(x) + ε may be more appropriate because by using (1), some form of centering for the y’s would be required. In fact, a transformation of y to center it at 0 is often done in the background of the BART algorithm (details given in the Appendix). We adopted (1) to be consistent with the definition used by most of the literature.

2.1.3 ∣. Sum of regression trees

We next focus on the sum of regression trees, $f(x)={\sum }_{j=1}^{m}g(x;T_j,M_j)$. We begin by using an example with m = 2 trees and p = 3 covariates. Suppose we were given the two trees in Figure 2. The resulting conditional mean of y given x is $E[y\mid x]={\sum }_{j=1}^{2}g(x;T_j,M_j)$. Consider the hypothetical data from n = 10 subjects given in Table 1. We can see that the quantity that is being ‘summed’ and eventually allocated to E(y∣x) is not the regression tree or tree structure, but the value that each j^th tree structure assigns to the subject. This is one way to think of a sum of regression trees. It allocates a sum of parameters μ_ji to E[y∣x] of the subject. Note that contrary to initial intuition, it is the sum of μ_ji’s that are allocated rather than the mean of the μ_ji’s. This is mainly because BART calculates each posterior draw of the regression tree function g(x; T_j, M_j) using a leave-one-out concept, which we shall elaborate shortly.

FIGURE 2.

Illustrating the sum of regression trees using a simple two regression tree example.

TABLE 1.

The values of $f(x)={\sum }_{j=1}^{2}g(x;T_j,M_j)$ from the regression trees in Figure 2.

Subject	y	x1	x2	x3	g(x; T1, M1)	g(x; T2, M2)	f(x)
1	y1	−182	235	−333	μ12	μ21	μ12 + μ21
2	y2	54	339	244	μ12	μ23	μ12 + μ23
3	y3	−106	−50	−682	μ11	μ21	μ11 + μ21
4	y4	−80	−62	−320	μ11	μ21	μ11 + μ21
5	y5	−123	198	−77	μ11	μ21	μ11 + μ21
6	y6	175	108	−46	μ13	μ21	μ13 + μ21
7	y7	−44	11	136	μ11	μ22	μ11 + μ22
8	y8	−131	−10	−70	μ11	μ21	μ11 + μ21
9	y9	−56	68	257	μ11	μ22	μ11 + μ22
10	y10	7	324	282	μ12	μ23	μ12 + μ23

Another way to view the concept of a sum of regression trees is to think of the regression trees in Figure 2 as ANOVA models (recall our single regression tree example). Then, the sum of regression trees for this simple example is just the following ANOVA model:

$$\begin{gathered} y=g(x;T_1,M_1)+g(x;T_2,M_2)+\epsilon \\ ={\mu }_{11}I\{x_1<100\}I\{x_2<200\}+{\mu }_{12}I\{x_1<100\}I\{x_2\ge 200\}+{\mu }_{13}I\{x_1\ge 100\} \\ +{\mu }_{21}I\{x_3<100\}+{\mu }_{22}I\{x_3\ge 100\}I\{x_2<200\}+{\mu }_{23}I\{x_3\ge 100\}I\{x_2\ge 200\}+\epsilon . \end{gathered}$$

Note that this does not imply that BART is an ANOVA model. This is because, in BART, each MCMC iteration produces a possibly different ANOVA model. These models are posterior draws from the target model distribution f(x).

Non-linearity of BART

From this simple example, we can see how BART handles non-linearity. Each single regression tree is a simple step-wise function or ANOVA model. When we sum regression trees together, we are actually summing together these ANOVA models or step-wise functions, and as a result, we eventually obtain a more complicated step-wise function which can approximate the non-linearities in the main effect and multiple-way interactions. It is this ability to handle non-linear main and multiple-way interaction effects that makes BART a flexible model. But unlike many flexible models, BART does not require the researcher to specify the main and multi-way interaction effects.

Prior distributions

In the examples above, we have taken the trees as a given, including which variables to split on, the splitting values, and the mean parameters at each terminal node. In practice, each g(x; T_j, M_j) is unknown. We therefore need prior distributions for these functions. Thus, we can also think of BART as a Bayesian model where the mean function itself is unknown. A major advantage of this approach is that uncertainty about both the functional form and the parameters will be accounted for in the posterior predictive distribution of y.

Before getting into the details of the prior distributions and MCMC algorithm, we will first walk through a simple example to build the intuition.

2.1.4 ∣. BART machinery: a visual perspective

In our simple example, we have three covariates x = (x₁, x₂, x₃) and a continuous outcome y. We will now run the BART MCMC algorithm with m = 4 regression trees for 5 iterations on this data set. At each iteration, we present the regression tree structures to illustrate how the algorithm works at each step. As mentioned previously, we typically would choose a large number for m, for example, 50. Here we are using just 4 trees to illustrate the key ideas.

The first step in the algorithm is initialization of the four regression trees to single root nodes (See “Initiation” in Figure 3). Since all four regression trees are single root nodes, the parameters initialized for these nodes would be ${\mu }_{ji}^{(0)}=\frac{\overline{y}}{m}=\frac{\overline{y}}{4}$.

FIGURE 3.

Initiation of BART to iteration 3 of the MCMC steps within BART with m = 4.

With this initialization in place, BART starts to draw the tree structures for each regression tree in the first MCMC iteration. Without loss of generality, let us start with determining (T₁, M₁), the first regression tree. This is possible because the ordering of the regression tree calculation does not matter when computing the posterior distribution. We first calculate $r_1=y-{\sum }_{j\ne 1}g(x;T_j,M_j)=y-[g(x;T_2,M_2)+g(x;T_3,M_3)+g(x;T_4,M_4)]=y-3\times \frac{\overline{y}}{4}$. Then a MH algorithm is used to determine the posterior draw of the tree structure, T₁ for this iteration. The basic idea of MH is to propose a new tree structure from T₁, call this $T_1^{\ast }$, and then calculate the probability of whether $T_1^{\ast }$ should be accepted, taking into consideration: $r_1\mid T_1^{\ast }$ (the likelihood of the residual given the new tree structure), r₁∣T₁ (the likelihood of the residual given the previous tree structure), the probability of observing $T_1^{\ast }$, the probability of observing T₁, the probability of moving from $T_1^{\ast }$ to T₁, and the probability of moving from T₁ to $T_1^{\ast }$. For example, if the new tree proposed $T_1^{\ast }$ consists of a single root node of x₂ < 100, then the likelihood of $r_1\mid T_1^{\ast }$ would be $N({\mu }_{11}^{(0)}I\{x_2<100\}+{\mu }_{12}^{(0)}\{x_2\ge 100\},{\sigma }^2)$ (recall the ANOVA representation of a regression tree) and the likelihood of $r_1\mid T_1$ would be $N({\mu }_{11}^{(0)},{\sigma }^2)$. More rigorous and technical details are provided in the Appendix. We describe the different types of moves from T₁ to $T_1^{\ast }$ in detail in the next subsection. If $T_1^{\ast }$ is accepted, T₁ is updated to become $T_1^{\ast }$ i.e. $T_1=T_1^{\ast }$. Else, nothing would be changed for T₁. From Figure 3, we can see that $T_1^{\ast }$ was not accepted in the first iteration so the tree structure remains as a single root node. The algorithm then updates M₁ based on the new updated regression structure for T₁, in our context draw ${\mu }_{11}^{(1)}$ from ${\mu }_{11}^{(1)}\mid T_1$, r₁, σ, and moves on to determine (T₂, M₂). Details of how to calculate ${\mu }_{ji}^{(b)}\mid T_j$, r_j, σ where b is the iteration index can be found in the Appendix.

To determine (T₂, M₂), again the algorithm calculates $r_2=y-{\sum }_{j\ne 2}g(x;T_j,M_j)=y-[g(x;T_1,M_1)+g(x;T_3,M_3)+g(x;T_4,M_4)]=y-({\mu }_{11}^{(1)}+2\times \frac{\overline{y}}{4})$, where ${\mu }_{11}^{(1)}$ is the updated parameter for regression tree 1. Similarly, MH is used to propose a new $T_2^{\ast }$ and r₂ is used to calculate the acceptance probability to decide whether $T_2^{\ast }$ should be accepted. Again, we see from Figure 3 that $T_2^{\ast }$ was not accepted. For (T₃, M₃), the MH iteration result is more interesting because the newly proposed $T_3^{\ast }$ was accepted and we can see from Figure 3 that a new tree structure was used for T₃ in Iteration 1. This new tree structure produces two terminal node parameters ${\mu }_{31}^{(1)}$ and ${\mu }_{32}^{(1)}$ which can be drawn from ${\mu }_{31}^{(b)}\mid T_3$, r₃, σ and ${\mu }_{32}^{(b)}\mid T_3$, r₃, σ respectively. As a result, when calculating r₄, this becomes $r_4=y-({\mu }_{11}^{(1)}+{\mu }_{12}^{(1)}+{\mu }_{31}^{(1)}I\{x_3<0.48\}+{\mu }_{32}^{(1)}I\{x_3\ge 0.48\}+{\mu }_{41}^{(1)})$. $T_4^{\ast }$ was not accepted and a single node was updated as the tree structure for (T₄, M₄). Once the regression tree draws are complete, BART then proceeds to draw the posterior distribution of σ², the variance of the error term in (1). More details will be given in the next subsection.

Figures 3 and 4 give the full iterations from initiation to iteration 5. From these figures we can see how the four regression trees grow and change from one iteration of the MCMC to another. This iterative process runs for a burn-in period (typically 100 to 1 000 iterations), before those draws are discarded, and then run for as long as needed to obtain a sufficient number of draws from the posterior distribution of f(x). After any full iteration in the MCMC algorithm, we have a full set of trees. We can therefore obtain a predicted value of y for any x of interest (simply by summing the terminal node μ_ji’s). By obtaining predictions across many iterations, we can easily obtain a 95% prediction interval. Another point to note is how shallow the regression trees are in Figures 3 and 4 with a maximum depth of 3. This is because the regression trees are heavily penalized (via the prior) to reduce the likelihood for a single tree to grow very deep. This concept (boosting) is similar to that found in the machine learning literature, where many weak models (learners) combined together can outperform a single strong model.

FIGURE 4.

Iterations 4 and 5 of the MCMC steps within BART with m = 4.

2.1.5 ∣. A rigorous perspective on the BART algorithm

Now that we have a visual understanding of how the BART algorithm works, we shall give a more rigorous explanation. First, we start with the prior distributions for BART. The prior distribution for (1) is p[(T₁, M₁), …, (T_m, M_m), σ]. The assumption is that {(T₁, M₁), …, (T_m, M_m)} and σ are independent and that (T₁, M₁), …, (T_m, M_m) are independent of each other. Then the prior distribution can be written as

$$\begin{gathered} p[(T_1,M_1),\dots ,(T_m,M_m),\sigma ]=p[(T_1,M_1),\dots ,(T_m,M_m)]p(\sigma ) \\ =\left[\prod _{j=1}^{m}p(T_j,M_j)\right]p(\sigma ) \\ =\left[\prod _{j=1}^{m}p(M_j\mid T_j)p(T_j)\right]p(\sigma ) \\ =\left[\prod _{j=1}^m\left\{\prod _{i=1}^{b_j}p({\mu }_{\mathit{ji}}\mid T_j)\right\}p(T_j)\right]p(\sigma ). \end{gathered}$$ (2)

For the third to fourth line in (2)recall that M_j = {μ_j1,…, μ_ibj} is the vector of terminal node parameters associated with T_j and each node parameter μ_ji is usually assumed to be independent of each other. Equation (2) implies that we need to specify priors μ_ji∣T_j, σ, and T_j. The priors for μ_ji∣_Tj and σ are usually given as ${\mu }_{ji}\mid T_j~N({\mu }_{\mu },{\sigma }_{\mu }^2)$ and ${\sigma }^2~IG\left(\frac{\nu }{2},\frac{\nu \lambda }{2}\right)$ respectively, where IG(a, b) is the Inverse-Gamma distribution with shape parameter a and rate parameter b. Note that the representation of the prior for σ² is a slight departure from most BART literature which often use the Inverse-Chisquare to represent the same prior distribution. We employed the Inverse-Gamma representation here because it is a more convenient form for the benefit of the reader.

The prior for p(T_j) is more interesting and can be specified using:

1. The probability that a node at depth d = 0, 1, … would split, which is given by $\frac{\alpha }{(1+d{)}^{\beta }}$. The parameter a ∈ {0, 1} controls how likely a node would split, with larger values increasing the likelihood of a split. The number of terminal nodes is controlled by parameter β > 0, with larger values of β reducing the number of terminal nodes. This aspect is important as this is the penalizing feature of BART which prevents BART from overfitting and allowing convergence of BART to the target function f(x)⁴².

2. The distribution used to select the covariate to split upon in an internal node. The default suggested distribution is the uniform distribution. Recent work by Ročková and van der Pas⁴³ and Linero⁴⁴ have argued that the uniform distribution does not promote variable selection and should be replaced if variable selection is desired.

3. The distribution used to select the cutoff point in an internal node once the covariate is selected. The default suggested distribution is the uniform distribution.

The setting of the other parameters for the BART priors is rather technical, so we refer interested readers to the Appendix.

The prior distribution would induce the posterior distribution:

$$p[(T_1,M_1),\dots ,(T_m,M_m),\sigma \mid Y]\propto p[Y\mid (T_1,M_1),\dots ,(T_m,M_m),\sigma ]\times p[(T_1,M_1),\dots ,(T_m,M_m),\sigma ]$$

which can be simplified into two major posterior draws using Gibbs sampling. First, draw m successive

$$p[(T_j,M_j)\mid T_{-j},M_{-j},y,\sigma ]$$ (3)

for j = 1,…, m, where T_−j and M_−j consist of all the tree structures and terminal nodes except for the j^th tree structure and terminal node; then, draw

$$p[\sigma \mid (T_1,M_1),\dots ,(T_m,M_m),y]$$ (4)

from $IG(\frac{\nu +n}{2},\{\nu \lambda +\sum [y-f(x){]}^2\}∕2)$.

To obtain a draw from (3), note that this distribution depends on (T_−j, M_−j, y, σ) through

$$r_j=y-\sum _{h\ne j}g(x;T_h,M_h),$$

the residuals of the m–1 regression sum of trees fit excluding the j^th tree (recall our visual example in the previous subsection). Thus (3) is equivalent to the posterior draw from a single regression tree r_j = g(x; T_j, M_j) + ε or

$$p[(T_j,M_j)\mid r_j,\sigma ].$$ (5)

We can obtain a draw from (5) by first integrating out M_j to obtain p(T_j∣r_j, σ). This is possible since a conjugate normal prior on μ_ji was employed. We draw p(T_j∣r_j, σ) using a MH algorithm where, first, we generate a candidate tree $T_j^{\ast }$ for the j^th tree with probability distribution $q(T_j,T_j^{\ast })$. The details of the new tree proposal are given in the next paragraph. We then accept $T_j^{\ast }$ with probability

$$\alpha (T_j,T_j^{\ast })=\text{min}\left\{1,\frac{q(T_j^{\ast },T_j)}{q(T_j,T_j^{\ast })}\frac{P(r_j\mid x,T_j^{\ast },M_j)}{P(r_j\mid x,T_j,M_j)}\frac{P(T_j^{\ast })}{P(T_j)}\right\}.$$ (6)

where $\frac{q(T_j^{\ast },T_j)}{q(T_j,T_j^{\ast })}$ is the ratio of the probability of how the previous tree moves to the new tree against the probability of how the new tree moves to the previous tree, $\frac{P(r_j\mid x,T_j^{\ast },M_j)}{P(r_j\mid x,T_j,M_j)}$ is the likelihood ratio of the new tree against the previous tree, and $\frac{P(T_j^{\ast })}{P(T_j)}$ is the ratio of the probability of the new tree against the previous tree.

A new tree $T_j^{\ast }$ can be proposed given the previous tree T_j using four local steps: (i) grow, where a terminal node is split into two new child nodes; (ii) prune, where two terminal child nodes immediately under the same non-terminal node are combined together such that their parent non-terminal node becomes a terminal node; (iii) swap, the splitting criteria of two non-terminal nodes are swapped; (iv) change, the splitting criteria of a single non-terminal node is changed. Once we have the draw of p(T_j∣r_j, σ), we then draw $p({\mu }_{ji}\mid T_j,r_j,\sigma )~N([{\sigma }_{\mu }^2\sum r_{ji}]∕[n_i{\sigma }_{\mu }^2+{\sigma }^2],[{\sigma }^2{\sigma }_{\mu }^2]∕[n_i{\sigma }_{\mu }^2+{\sigma }^2])$, where r_ji is the subset of elements in r_j allocated to the terminal node parameter μ_ji and n_i is the number of r_ji’s allocated to μ_ji. We derive p(μ_ji∣_Tj, r_j, σ), (4), and (6) for the grow and prune steps as an example in the Appendix. Other possible moves have also been proposed in literature. Interested readers can refer to the following papers^45,46,47,48.

2.2 ∣. Binary outcomes

For binary outcomes, BART can be extended using a probit model. Specifically, we may write

$$\mathit{P}(y=1\mid x)=\Phi [f(x)]$$ (7)

where Φ[.] is the cumulative distribution function of a standard normal distribution. With such a setup, only priors for (T₁, M₁),…, (T_m, M_m) are needed. The same decomposition in (2) without σ can be employed and the similar prior specifications for μ_ji∣T_j and T_j can be used. An important note that some readers may not notice is that such a setup in (7) shrinks f(x) toward 0 and hence shrinking Φ[f(x)] toward 0.5. This may not be what researchers want nor anticipate and hence, if it is more desirable to shrink Φ[f(x)] toward a value p₀ ∈ (0, 1) that is not 0.5, the strategy is to replace f(x) in (7) with f_c = f(x) + c where the offset c = Φ⁻¹[p₀]. The setup of the remaining parameters are slightly different from that of continuous outcomes and we describe the details in the Appendix.

To estimate the posterior distribution, data augmentation⁴⁹ can be used. We assume that y = I(z > 0) where z is a latent variable drawn as follows:

$$\begin{gathered} z~N_{(-\infty ,0)}[f(x),1]\text{if}y=0, \\ z~N_{(0,\infty )}[f(x),1]\text{if}y=1 \end{gathered}$$

with N_(a,b))[μ, σ²] being a truncated normal distribution with mean μ and variance σ² truncated at (a, b). Next, we can treat z as the continuous outcome for a BART model with

$$z=f(x)+\epsilon$$ (8)

where ε ~ N(0, 1) because we employed a probit link. The usual posterior estimation for a continuous outcome BART with σ ≡ 1 can now be employed on (8) for one iteration in the MCMC. The updated f(x) can then be used to draw a new z and this new z can be used to draw another iteration of f(x). The process can then be repeated until convergence.

3 ∣. ILLUSTRATING THE PERFORMANCE OF BART

3.1 ∣. Posterior performance via synthetic data

We generated a synthetic data set with p = 3, n = 1000 and, the true model for y is

$$y=0.5+0.1x_1+0.3x_2^2+0.7\sin (x_3)+0.2x_1{x}_2+0.9\sqrt{\mid x_1{x}_3\mid }+0.4\text{exp}(x_2{x}_3)+0.8\log (\mid x_1{x}_2{x}_3\mid )+\epsilon$$

with each x_p ~ N(0, 1) and ε ~ N(0, 2). The goal is to demonstrate that BART can predict y’s effectively even in complex, non-linear models, and also properly account for prediction uncertainty. We will compare results of BART against that of a parametric Bayesian linear regression (BLR) model. To this end, we randomly selected 970 samples as the training set and then use the remaining 30 samples as the testing set. Note that in practice, we commonly see a 60(training)-40(testing) or a 70(training)-30(testing) split being used although more rigorous splits have been suggested⁵⁰. We employed a different split because a graphical presentation of 300 or 400 results would not be feasible. We also varied the number of trees used by BART to illustrate how varying m affects the performance of BART. We plotted the point estimate and 95% credible interval of the 30 randomly selected testing data points and compared them with their true values in Figure 5. The code to implement this simulation can be found in our Supporting Information.

FIGURE 5.

Posterior mean and 95% credible interval of Bayesian linear regression (BLR) and BART with m = 1, 50, 100, 150, 200 for 30 randomly selected testing set outcomes. n = 1000, black=true value, red=model estimates.

We can see from Figure 5 that most of the point estimates of BLR were far away from their true values and many of the true values were not covered by the 95% credible interval. For BART with a single tree, although the true values were mostly covered by the 95% credible interval, the point estimates were far from their true values. When we increased the number of trees to 50 in BART, we see a significant improvement in terms of bias (closeness to the true values) compared to both BLR and BART with m = 1. In addition, we see a narrowing of the 95% intervals. As we increase the number of trees, the point estimate and 95% intervals stabilize. In other words, we might see a big difference between m = 1 and m = 50, and virtually no difference between m = 200 and m = 20 000. In practice, the idea is to choose a large enough value for m so that BART approximates the results that would have been obtained if more trees were used. Although as we increase m, the computation efficiency would decrease, the fact that each regression tree is a weak learner still allows BART to finish computation within a reasonable amount of time. Hence, Chipman et al.¹ suggested m = 200 originally but later researchers agree that m = 50 is often adequate⁴¹. Alternatively, one could determine a sufficiently large m using cross validation¹, but that would be computationally expensive.

3.2 ∣. Predicting the Standardized Hospitalization Ratio from the 2013 Centers for Medicare and Medicaid Services Dialysis Facility Compare data set

We next present an example to demonstrate how BART can be applied to a data set to improve prediction over the usual multiple linear regression (MLR) model. The 2013 Centers for Medicare and Medicaid Services Dialysis Facility Compare data set contains information regarding 105 quality measures and 21 facility characteristics of all dialysis facilities in the US, including US territories. This data set is available publicly. We provide the data set we downloaded and code in our Supporting Information. We are interested in finding a model that can better predict the standardized hospitalization ratio (SHR). This quantity is important because a large portion of dialysis cost for End Stage Renal Disease (ESRD) patients can be attributed to patient hospitalizations.

Table 2 shows some descriptive statistics for this data set. SHR was adjusted for a patient’s age, sex, duration of ESRD, comorbidities, and body mass index at ESRD incidence. We removed 463 facilities (7%) with missing SHR values because of small patient numbers. We also removed peritoneal dialysis (PD) removal greater than 1.7 Kt/V because of the high proportion of missingness (80%). We combined pediatric hemodialysis (HD) removal greater than 1.2 Kt/V with adult HD removal greater than 1.2 Kt/V because most facilities (92%) do not provide pediatric HD. We re-categorized the chain names to “Davita”, “Fresenius Medical Care (FMC)”, “Independent”, “Medium”, and “Small”. “Medium” consists of chains with 100-500 facilities while “Small” are chains with less than 100 facilities. To estimate patient volume, we used the maximum of the number of patients reported by each quality measure group: Urea Reduction Ratio (URR), HD, PD, Hemoglobin (HGB), Vascular Access, SHR, SMR, STR, Hypercalcemia (HCAL), and Serum phosphorus (SP). We also logarithm-transformed (log) SHR, SMR, and STR so that the theoretical range for these log standardized measures will be −∞ to ∞.

TABLE 2.

Descriptive statistics of dialysis facility characteristics and quality measures (n=5,774).

Parameters (% missing)	Mean (s.d)/Frequency (%)	Parameters (% missing)	Mean (s.d)/Frequency (%)
Arterial Venous Fistula (3)	63.27 (11.22)	Number of stations	18.18 (8.27)
Avg. Hemoglobin<10.0 g/dL (5)	12.86 (10.32)	Serum P. (2)	28.52 (5.13)
Chain name:		Shift after 5pm?
Davita	1,812 (31)	Yes	1,097 (19)
FMC	1,760 (30)	No	4,677 (81)
Independent	820 (14)	SHR	1.00 (0.31)
Medium	740 (13)	SMR (2)	1.02 (0.29)
Small	642 (12)	STR (7)	1.01 (0.54)
Patient volume*	100.07 (60.93)	Type:
Facility Age (years)	14.47 (9.81)	All (HD, Home HD, & PD)	1,443 (25)
For profit?		HD & PD	1,897 (33)
Yes	4,967 (86)	HD & Home HD	103 (2)
No	806 (14)	HD alone	2,331 (40)
HD≥1.2 Kt/V (4)	88.52 (9.85)	URR≥65% (7)	98.77 (3.04)
Hypercalcemia (3)	2.37 (3.20)	Vas. Catheter>90 days (3)	10.74 (6.66)

^*Estimated

For our analysis, we used the log-transformed SHR as the outcome and the variables in Table 2 as the predictors. We used the root mean squared error (RMSE) of a 10-fold cross-validation to compare the prediction performance from multiple linear regression (MLR), Random Forest (RF), and BART. For RF and BART, we used the default settings from the R packages randomForest and BayesTree respectively. The 10 RMSEs produced by each method from the 10-fold cross validation is provided in Figure 6. It is clear from this figure that BART and RF produce very similar prediction performances and are better compared to MLR. The mean of these 10 values also suggested a similar picture with MLR producing a mean of 0.24 while RF and BART produced a mean of 0.23.

FIGURE 6.

Root mean squared error (RMSE; y-axis) for the 10-fold cross-validation of multiple linear regression (MLR), random forest, and Bayesian additive regression trees of log transformed standardized hospitalization ratio (SHR). x-axis indicates the RMSE for the xth fold.

3.3 ∣. Predicting left turn stops at an intersection

We next present another illustration of BART on a real example, but this time with a binary outcome. In Tan et al.⁵¹, the authors were interested in predicting whether a human driven vehicle would stop at an intersection before making a left turn. Left turns are important in countries with right side driving because most vehicle conflicts including crashes at intersections occur during left turns. Accurate predictions about whether a human driven vehicle would stop before executing a left turn could help driverless vehicles improve decision making at intersections. More details about this data set can be found in Tan et al.⁵¹. In brief, the data comes from the Integrated Vehicle Based Safety System (IVBSS) study conducted by Sayer et al.⁵². This study collected driving data from 108 licensed drivers in Michigan between April 2009 and April 2010. Each driver drove one of the 16 research vehicles fitted with various recording devices to capture the vehicle dynamics while the subject is driving on public roads for 12 days.

For this example, we focused on the vehicle speed at 50, 51, …, 56 meters away from the center of an intersection and whether we could utilize this vehicle dynamic to predict whether the driver would stop (vehicle speed < 1m/s) in the future (0, 1, …, 49 meters from the center of an intersection). Similar to Tan et al.⁵¹, we performed a principal components analysis (PCA) on the vehicle speeds (50, 51,…, 56 meters) and took the first 3 principal components (PCs) as the predictors. Tan et al.⁵¹ took the first 3 PCs as predictors because they found that using the first 3 PCs as predictors improved prediction performance compared to using the original vehicle speeds. In addition, adding more PCs beyond the first 3 did not improve prediction performance. We ran a 10-fold cross-validation on this data set and compared the binary prediction results of logistic regression, RF, and BART. Since the outcome of interest for this data set was binary, we used the area under the receiver operating curve (AUC) to determine the prediction performance instead of the RMSE, which is more suited for continuous outcomes.

Figure 7 shows the results of the 10 AUCs produced by each method from the 10-fold cross validation. BART performed better than either logistic regression or RF in predicting whether the human-driven vehicle would stop in the future at an intersection before making a left turn. This is also evident from the mean of the 10 cross validation AUC values produced by each method. BART produced a mean of 0.81 compared to 0.79 from logistic regression, and 0.78 from RF.

FIGURE 7.

Area under the receiver operating characteristic curve (AUC; y-axis) for the 10-fold cross-validation of logistic regression, random forest, and Bayesian additive regression trees of left turn stop probabilities at an intersection. x-axis indicates the AUC for the xth fold.

4 ∣. GENERAL BART MODEL

Recently researchers have extended or generalized BART to a wider variety of settings, including clustered data, spatial data, semiparametric models, and to situations where more flexible distributions for the error term is needed. Taking a closer look at these recent extensions, we found a common and unifying theme. It enables the proposed models to be fitted without having to do extensive re-derivation for the posterior algorithm described in Section 2. By unifying the models in these various papers into one single approach, we believe it will aid the application of BART to a wider variety of situations. In addition, we hope the discussion of our approach would be able to spark inspiration for further, possibly more complicated extensions of BART.

To set up our General BART model, suppose once again that we have a continuous outcome y and p covariates x = {x₁,…, x_p}. Suppose also that we have another set of q covariates w = {w₁,…,w_q}, such that no two columns in x and w are the same. Then, we can extend (1) as follows:

$$y=f(x)+h(w,\mathrm{\Theta })+\epsilon$$ (9)

where h(.) is a function that works on w using parameter Θ, and ε ~ G(Σ) can be any distribution with parameter Σ.

Assuming that {(T₁, M₁), …, (T_m, M_m)}, Θ, and £ are independent, the prior distribution for (9) would be p[(T₁, M₁),…, (T_m, M_m)]p(Θ)p(Σ). Assuming again that the (T_j, M_j)’s are independent of each other, p[(T₁, M₁),…, (T_m, M_m)] can be decomposed into ${\prod }_{j=1}^m\left\{{\prod }_{i=1}^{b_j}p({\mu }_{ji}\mid T_j)\right\}p(T_j)$. The priors needed are thus p(μ_ji∣T_j), p(T_j), p(Θ), and p(Σ). Note that it is possible to model Θ and Σ jointly so that the prior distribution becomes ${\prod }_{j=1}^m\left\{{\prod }_{i=1}^{b_j}p({\mu }_{ji}\mid T_j)\right\}p(T_j)p(\mathrm{\Theta },\Sigma )$ instead. We shall see this in Example 4.4.

To obtain the posterior distribution of p[(T₁, M₁),…, (T_m, M_m), Θ, Σ∣y], Gibbs sampling can be used. For P[(T₁, M₁),…, (T_m, M_m)∣Θ, Σ, y], this can be seen as drawing from the following model

$$\tilde{y}=f(x)+\epsilon$$ (10)

where $\tilde{y}=y-h(w,\mathrm{\Theta })$, which is just a BART model with a modified outcome $\tilde{y}$. Hence, the BART algorithm presented in Section 2 can be used to draw {(T₁, M₁),…, (T_m, M_m)}. Similarly, p[Θ∣(T₁, M₁),…, (T_m, M_m), Σ, y] can be obtained by drawing from the model

$$y^{\prime }=h(w,\mathrm{\Theta })+\epsilon$$ (11)

where y′ = y – f(x). This posterior draw depends on the function h(.) being used as well as the prior distribution specified for Θ. As there are many possibilities where we can set up h(.) and Θ, we shall not discuss the specifics here. The examples we present in the subsequent subsections will highlight a few of these possibilities we have seen in the literature thus far. Finally, drawing from p[Σ∣(T₁, M₁),…, (T_m, M_m), Θ, y] is just drawing from the model

$$r=y-f(x)-h(w,\mathrm{\Theta })=\epsilon .$$ (12)

Again, many possibilities are available for setting up the prior distribution for Σ and hence the distributional assumption for ε. The default is usually ε ~ N(0, σ²) where $\Sigma ={\sigma }^2~IG\left(\frac{\nu }{2},\frac{\nu \lambda }{2}\right)$. Example 4.4 shows a plausible alternative. Iterating through these Gibbs steps will give us the posterior draw of p[(T₁, M₁),…, (T_m, M_m), Θ, Σ∣y].

For binary outcomes, the probit link can once again be used where

$$\mathit{P}[y=1\mid x]=\Phi [f(x)+h(w,\mathrm{\Theta })].$$

Under this framework, we will only need priors for p[(T₁, M₁),…, (T_m, M_m)] and p(Θ). p[(T₁, M₁),…, (T_m, M_m)] can be decomposed once again into ${\prod }_{j=1}^m\left\{{\prod }_{i=1}^{b_j}p({\mu }_{ji}\mid T_j)\right\}p(T_j)$ if we are willing to assume that the m trees are independent of one another, and data augmentation⁵³ can be used obtain the posterior distribution. Again with the assumption that y = I(z > 0), we can draw

$$\begin{gathered} z~N_{(-\infty ,0)}(f(x)+h(w,\mathrm{\Theta }),1)\text{if}y=0, \\ z~N_{(0,\infty )}(f(x)+h(w,\mathrm{\Theta }),1)\text{if}y=1 \end{gathered}$$

and then treat z as the outcome for the model in (9). This implies ε ~ N(0, 1) and we can apply the Gibbs sampling procedure we described for continuous outcomes using z instead of y with Σ = σ = 1. Iterating through the latent draws and Gibbs steps will produce the posterior distribution that we require.

With the general framework and model for BART in place, we are now equipped to consider how Zeldow et al.³⁷, Tan et al.³⁸, Zhang et al.³⁹, and George et al.⁴⁰ extended BART to solve their research problems in the next four subsections.

4.1 ∣. Semiparametric BART

The semiparametric BART approach was first developed by Zeldow et al.³⁷. Their work was motivated by observational studies where there are many confounders, an exposure, and a small number of variables that are of interest as exposure-covariate interaction (heterogenity of treatment effect). While a flexible model such as BART could handle the potential complex confounding and main effects relationship with the outcome, there is some loss of interpretability relative to a parametric model. On the other hand, fully parametric models rely on strong assumptions.

Therefore, their idea was to have a semiparametric model, where the exposure and exposure-covariate interactions had parametric specification, and the nuisance or confounder variables were modeled nonparametrically (using BART). In its simplest form, we can reconstruct this idea under the framework of (9) as follows. The x covariates are the confounders that have to be controlled for, but are not of primary interest. The variables w include the exposure and covariates that are of substantive interest as possible effect modifiers. Then,

$$h(w,\mathrm{\Theta })={\theta }_0+{\theta }_1{w}_1+\dots +{\theta }_q{w}_q$$

where w = {w₁,…, w_q}, Θ = {θ₀,…, θ_q}, and ε ~ N(0, σ²) with Σ = σ. Prior distributions for μ_ji∣T_j, T_j, and σ² follow the usual distributions we use for BART while Θ ~ MVN (β, Ω). Posterior estimation follows the procedure we described in Section 4 using Gibbs Sampling. For (10) and (12), since they suggested using the default BART priors, the usual BART mechanisms can be applied to obtain the posterior draws. For (11), Θ ~ MVN(β, Ω) implies that we can treat this as the usual BLR and standard Bayesian methods could be used to obtain the posterior draw for Θ. The framework for binary outcomes follows easily using the data augmentation step we describe in Section 4.

4.2 ∣. Random intercept BART for correlated outcomes

Random intercept BART (riBART) was proposed by Tan et al.³⁸ as a method to handle correlated continuous or binary outcomes with correlated binary outcomes as the main focus. In this work, the authors wanted to predict whether a human driven vehicle would stop at an intersection before making a left turn. They used a data set that contained about 100 drivers with each driver making numerous turns. In a preliminary analysis, they found that BART produced better prediction performance. However, BART was designed for independent outcomes and not correlated binary outcomes. Hence, they proposed to extend BART to handle the correlated binary outcomes in their data set by using a random intercept (extensions to random intercept and slope is straight forward).

To re-formulate riBART under the framework of (9), we set h(w, Θ) = wa, where Θ = (a, τ) and

$$w=\left[\begin{array}{cccc}\hfill 1\hfill & \hfill 0\hfill & \hfill \dots \hfill & \hfill 0\hfill \\ \hfill ⋮\hfill & \hfill ⋮\hfill & \hfill \ddots \hfill & \hfill ⋮\hfill \\ \hfill 1\hfill & \hfill 0\hfill & \hfill \dots \hfill & \hfill 0\hfill \\ \hfill 0\hfill & \hfill 1\hfill & \hfill \dots \hfill & \hfill 0\hfill \\ \hfill ⋮\hfill & \hfill ⋮\hfill & \hfill \ddots \hfill & \hfill ⋮\hfill \\ \hfill 0\hfill & \hfill 1\hfill & \hfill \dots \hfill & \hfill 0\hfill \\ \hfill ⋮\hfill & \hfill ⋮\hfill & \hfill ⋮\hfill & \hfill ⋮\hfill \\ \hfill 0\hfill & \hfill 0\hfill & \hfill \dots \hfill & \hfill 1\hfill \\ \hfill ⋮\hfill & \hfill ⋮\hfill & \hfill \ddots \hfill & \hfill ⋮\hfill \\ \hfill 0\hfill & \hfill 0\hfill & \hfill \dots \hfill & \hfill 1\hfill \end{array}\right]$$

i. e. w is a n × K matrix, where k indexes the subjects, $n={\sum }_{k=1}^K{n}_k$, with the first column being a vector of $(1_{n_1}{0}_{n-n_1}{)}^T$, the second column $(0_{n_1}{1}_{n_2}{0}_{n-n_1-n_2}{)}^T$, and so on until the last column $(0_{n-{\sum }_{k=1}^{K-1}{n}_k}{1}_{n_K}{)}^T$, where 1_i and 0_j are row vectors of 1s and 0s with size i and j respectively. Let a = {a₁,…, a_K} where a∣τ² ~ N(0, τ²), ε ~ N(0, σ²) where Σ = σ, and specify the usual BART priors for σ, μ_ji∣_Tj, and T_j. a and ε are assumed to be independent. A simple prior of τ² ~ IG(1, 1) could be used although more robust or complicated priors are possible. Posterior estimation and binary outcomes then follow the procedure described in Section 4 easily.

4.3 ∣. Spatially-adjusted BART for a statistical matching problem

The spatially-adjusted BART approach of Zhang et al.³⁹ was proposed to handle statistical matching problems⁵⁴ that occur in surveys. In statistical matching problems, inference is desired for the relationship between two different variables collected from two different studies. For example, survey A may collect information on income but survey B collects information on health status. Often, subjects in both surveys do not overlap. The relationship between income and health status is then desired. To solve this, Zhang et al.³⁹ proposed to use the geographical information or information aggregated at the geographical level to impute the missing variables. They set up two different models for surveys A and B separately but with a shared spatial random effects which takes into account the geographical location of each individual. This spatial random effects allow correlation among subjects living in adjacent geographical locations to be employed to improve inference. Although their approach is specifically set up to handle statistical matching problems, a simpler version can be formulated to handle more general spatial data.

To do so, we once again use the general BART model with w and a being the same setup as Section 4.2 but with

$$a\mid \rho ,{\delta }^2~N(0,{\delta }^2(H-\rho C{)}^{-1})$$ (13)

where C = c_kl is a K × K adjacency matrix, l = 1,…, K, with c_kl = 1 if group k and group l are (spatial) neighbors for k ≠ l; c_kl = 0 otherwise; and c_kl = 0 if k = l. H is a diagonal K × K matrix with diagonals $h_k={\sum }_{l=1}^K{c}_{kl}$, ρ is a parameter with range (−1, 1), and δ² is the variance component for (13). The parameters ρ and δ² are prespecified. Finally, ε ~ N(0, σ²) and (9) is completed by placing the usual BART priors for σ, μ_ji∣T_j, and T_j. Posterior draws again follow the procedures we outline in Section 4.

4.4 ∣. Dirichlet Process Mixture BART

Similar to how BART can be used as a prior for unknown function f(·), Dirichlet process (DP) priors are often used as priors for unknown distributions⁵⁵. We use notation D ~ DP(D₀, α) if a random distribution D follows a DP prior with base measure D₀ and concentration parameter α.

For the General BART Model in (9), it is the distribution of ε that is unknown. The following DP Mixture BART was proposed by George et al⁴⁰:

$$\begin{gathered} {\epsilon }_k\stackrel{i.i.d.}{~}N(a_k,{\sigma }_k^2), \\ (a_k,{\sigma }_k^2)~D, \\ \mathit{D}~DP({\mathit{D}}_0,\mathit{\propto }). \end{gathered}$$

Although the assumption of ${\epsilon }_k\stackrel{i.i.d.}{~}N(a_k,{\sigma }_k^2)$ suggests that each subject will have their own mean and variance for the error term, draws from a DP are discrete with probability 1. Therefore, the number of unique values for (a_k, ${\sigma }_k^2$) at any step in the posterior algorithm will be much less than n. Essentially, the distribution of ε can be viewed as a mixture of normals, where the number of mixture components is not specified. Lower values of the concentration parameter α favor fewer mixture components. Therefore, it is common to specify a prior for α that puts more weight on smaller values (i.e., our prior guess is that we will not need many mixture components).

Viewing DPMBART as a form of (9), we have h(w, Θ) = wa where w and a have the same structure as riBART and $p(\mathrm{\Theta },\Sigma )=(a_k,{\sigma }_k^2)$. Note here that we are no longer assuming that a_k and ε_k are independent unlike in some of our previous examples.

DPMBART requires specification of D₀ and a prior for α. For D₀, the commonly employed form is p(μ, σ∣ν, λ, μ₀, h₀) = p(σ∣ν, λ)p(μ∣σ, μ₀, h₀). George et al.⁴⁰ specified their priors as

$${\sigma }^2~\mathit{IG}\left(\frac{v}{2},\frac{v\lambda }{2}\right);\mu \mid \sigma ~N\left({\mu }_0,\frac{{\sigma }^2}{h_0}\right).$$

The parameter ν is set at 10 to make the spread of error for a single component h tighter. The parameter λ is chosen using the idea from how λ is determined in BART with the quantile set at 0.95 instead of 0.9 (See Appendix for how λ is determined in BART). For μ₀, because DPMBART subtracts $\overline{y}$ from y, μ₀ = 0. For h₀, the residuals of a multiple linear regression fit is used to place μ into the range of these residuals, r. The marginal distribution of is $\mu ~\frac{\sqrt{\lambda }}{\sqrt{h_0}}{t}_{\nu }$, where t_ν is a t distribution with ν degrees of freedom. Let h_s be the scaling for μ. Given h_s = 10, h₀ can be chosen by solving

$$\text{max}\mid r_h\mid =h_s\frac{\sqrt{\lambda }}{\sqrt{h_0}}.$$

For α, the prior used by DPMBART is the same as in Section 2.5 of Rossi⁵⁶, where the idea is to relate α to the number of unique components in (a_k, ${\sigma }_k^2$).

The posterior draw for DPMBART follows most of the ideas discussed in General BART where first, the idea of (10) is used to draw (T₁, M₁),…, (T_m, M_m)∣a_k, ${\sigma }_k^2$. The slight difference is to view this as a weighted BART draw with ${\epsilon }_k\stackrel{i.i.d.}{~}N(0,w_k{\sigma }^2)$. The second draw, (a_k, σ²)∣(T₁, M₁),…, (T_m, M_m) follows (11) which can be solved by using draws (a) and (b) of the algorithm in Section 1.3.3 of Escobar and West⁵⁷. The final draw is α∣(a_k, σ²). This is obtained by putting α on a grid and using Bayes’ theorem with $p(\alpha \mid (a_k,{\sigma }_k^2))=p(\alpha \mid H)\propto p(H\mid \alpha )p(\alpha )$ where H is the number of unique (a_k, ${\sigma }_k^2$)’s. Another way of generating these posterior draws (DPM portion) can be found in Neal⁵⁸.

5 ∣. DISCUSSION

In this tutorial, we walked through the BART model and algorithm in detail, and presented a general model based on recent extensions. We believe this is important because of the growing use of BART in research applications. By clarifying the various components of BART, we hope that researchers will be more comfortable using BART in practice.

Despite the success of BART, there are a number of papers that point out some of its limitations and propose modifications. One such limitation is the suboptimal performance of BART with high-dimensional data due to the use of the uniform prior to select the covariate to be split upon in the internal nodes. A solution is to allow researchers to place different prior probabilities on each covariate⁵⁹. More sophisticated solutions include using a DP prior for selecting covariates⁴⁴ or using a spike-and-slab prior⁶⁰. Other solutions have also been proposed^61,41. Another commonly addressed limitation is the computational speed of BART. Due to the many MH steps that BART requires, computation speed of BART can often be slow, especially when the sample size n and/or the number of covriates p is large. One direction is to parallelize the computational steps in BART, for example, Pratola et al.⁶², Kapelner and Bleich⁵⁹, and Entezari et al.⁶³ The other direction is to improve the efficiency of the MH steps which leads to the reduction in the number of trees needed. Notable examples include Lakshminarayanan et al.⁶⁴, where particle Gibbs sampling was used to propose the tree structure T_j’s; Entezari et al.⁶³, where likelihood inflated sampling was used to calculate the MH steps, and more recently He et al.⁴⁸, where they proposed to use a different tree-growing algorithm which grows the tree from scratch (root node) at each iteration. Other less discussed issues with BART include the problem of under estimation of the uncertainty of BART caused by inefficient mixing when the true variation is small⁴⁷, inability of BART to handle smooth functions⁶⁵, and inclusion of many spurious interactions when the number of covariates is large⁶⁶. Finally, the posterior concentration properties of BART have also been discussed recently by Ročková and van der Pas⁴³, Ročková and Saha⁴², and Linero and Yang⁶⁵. These works provide theoretical proof of why BART has been successful in many data applications we have seen thus far.

A second component we focused on was how we can extend BART using a very simple idea without having to re-write the whole MCMC algorithm to draw the regression trees. We term this framework General BART. Special cases of this model have already been used by various authors to extend BART to a variety settings. By unifying these methods under a single framework and showing how these methods are related to the General BART model, we hope to provide researchers with the start of a roadmap for new extensions. For example, researchers working with longitudinal data may want a more flexible modeling portion for the random effects and hence may want to model h(w, Θ) as BART. Another possibility is to combine the ideas in Examples 4.1, 4.2, and 4.4, i.e. correlated outcomes with an interpretable linear model portion and robust error assumptions.

We do note that the critical component of our General BART framework is re-writing the model in such a way that the MCMC draw of the regression trees can be done separately from the rest of the model. In situations where this is not possible, re-writing of the MCMC procedure for the regression trees may be needed. An example of this would occur if, rather than mapping the outcome to a parameter at the terminal node of a regression tree, it is mapped to a regression model. However, we feel that the general BART model is flexible enough to handle many of the extensions that might be of interest to researchers.