Bayesian Additive Regression Trees using Bayesian model averaging

Abstract

Bayesian Additive Regression Trees (BART) is a statistical sum of trees model. It can be considered a Bayesian version of machine learning tree ensemble methods where the individual trees are the base learners. However, for datasets where the number of variables p is large the algorithm can become inefficient and computationally expensive. Another method which is popular for high-dimensional data is random forests, a machine learning algorithm which grows trees using a greedy search for the best split points. However, its default implementation does not produce probabilistic estimates or predictions. We propose an alternative fitting algorithm for BART called BART-BMA, which uses Bayesian model averaging and a greedy search algorithm to obtain a posterior distribution more efficiently than BART for datasets with large p. BART-BMA incorporates elements of both BART and random forests to offer a model-based algorithm which can deal with high-dimensional data. We have found that BART-BMA can be run in a reasonable time on a standard laptop for the “small n large p” scenario which is common in many areas of bioinformatics. We showcase this method using simulated data and data from two real proteomic experiments, one to distinguish between patients with cardiovascular disease and controls and another to classify aggressive from non-aggressive prostate cancer. We compare our results to their main competitors. Open source code written in R and Rcpp to run BART-BMA can be found at: https://github.com/BelindaHernandez/BART-BMA.git.

Appendices

Appendix A: BART full conditional distribution $p(R_j|X,T_j,\sigma ^{-2})$

Using the forms of $p(\mu _{ij})$ and $p(\sigma ^{-2})$ described in Sect. 2.2.2 gives rise to the following full conditional distribution of the partial residuals for the BART model:

$$\begin{aligned} p(R_j|X,T_j,\sigma ^{-2})\propto & {} \prod _{i=1}^b\left( n_i\sigma ^{-2} +\left( \frac{0.5}{e\sqrt{m}}\right) ^{-2}\right) ^{-\frac{1}{2}}{\sigma ^{-2}}^{\frac{n+\nu }{2}-1}\nonumber \\&\times \exp \left( -\frac{\sigma ^{-2}}{2}\left( \sum _{\iota =1}^{n_i} R_{\iota ij}^2 +\nu \lambda \right) \right) \nonumber \\&\exp \left( \frac{n_i^2\bar{R}_{i j}^2\sigma ^{-4}}{2\left( n_i\sigma ^{-2}+\left( \frac{0.5}{e\sqrt{m}}\right) ^{-2}\right) } \right) , \end{aligned}$$

(11)

where $n_i$ is the number of observations in terminal node i of tree j and $\bar{R}_{ij}$ is the mean of the partial residuals $R_j$ for terminal node i in tree j.

Appendix B: BART for classification

For binary classification, Chipman et al. (2010) follows the latent variable probit approach of Albert and Chib (1993). Latent variables $Z_k$ are introduced so that

$$\begin{aligned} Y_k={\left\{ \begin{array}{ll} 1 &{}\text {if }Z_k>0 \\ 0 &{}\text {otherwise} . \end{array}\right. } \end{aligned}$$

The sum of trees prior is then placed on the $Z_k$ so that $Z_k \sim N\left( \sum _{j=1}^m g(x_k;T_j,M_j),1 \right)$. It follows that $Y_k$ is Bernoulli with

$$\begin{aligned} P(Y_k=1|x_k) = \Phi \left( \sum _{j=1}^m g(x_k;T_j,M_j)\right) , \end{aligned}$$

(12)

where $\Phi$ is the standard normal cumulative distribution function (CDF), used here as the link function. Note that there is no residual variance parameter $\sigma ^2$ in the classification version of the model.

Using the same prior distribution structure as in Sect. 2.2.2, the full posterior distribution of this version of the model is:

$$\begin{aligned} \begin{aligned}&p(T,M,Z|X,Y) \propto p(Y|Z) p(Z|X,T,M)\\&\quad \times \left[ \prod _{j}\prod _{i}p(\mu _{ij}|T_j)p(T_j)\right] , \end{aligned} \end{aligned}$$

(13)

where the top level of the likelihood (i.e., the first term on the right-hand side) is a deterministic function of the latent variables. The conditional prior distributions of the terminal node parameters $\mu _{ij}|T_j$ are set exactly as described in Sect. 2.2.2 except that $\sigma _0=\frac{3}{e\sqrt{m}}$ instead of $\sigma _0=\frac{0.5}{e\sqrt{m}}$. This is in order to assign high prior probability to the interval $(\Phi [-3],\Phi [3])$ which corresponds to the 0.1 and $99.9\%$ quantiles of the normal CDF.

The fitting algorithm proposed by Chipman et al. (2010) for the classification model is nearly identical to that of their standard algorithm. The only difference is that the latent variables $Z_k$ introduce an extra step in the Gibbs algorithm. The full conditional distributions of $Z_k|\ldots$ are:

$$\begin{aligned} Z_k|\ldots \sim {\left\{ \begin{array}{ll} \min {\left[ N \left( \sum _j g(x_k;T_j,M_j),1 \right) ,0\right] } &{}\text {if }Y_k=1, \\ \max {\left[ N \left( \sum _j g(x_k;T_j,M_j),1 \right) , 0\right] } &{}\text {if }Y_k=0. \end{array}\right. } \end{aligned}$$

(14)

The partial residuals in the tree updates are of course now based on the latent variables $Z_k$ for updating of individual trees.

Appendix C: BART-BMA algorithm overview

Appendix D: BART-BMA post hoc Gibbs sampler

In order to provide credible intervals for the point predictions, $\hat{Y}$, provided by BART-BMA, we run a post hoc Gibbs sampler. For each sum of trees model $\mathcal {T}_\ell$ in Occam’s window, a separate chain in the MCMC algorithm is run. For each model $\mathcal {T}_\ell$, each terminal node parameter $\mu _{ij}$ in each tree $T_j$ is then updated followed by an update of $\sigma ^{2}$. The details of the updates for the full conditional of $p(\mu _{ij}|T_j,R_j,\sigma ^2)$ and of $p(\sigma ^2)$ are explained in further detail in the following sections. The Gibbs sampler yields credible and prediction intervals for each set of sum of trees models accepted by BART-BMA along with the updates for $\sigma ^{-2}$ for each set of trees accepted in the final BART-BMA model. The final simulated sample from the overall posterior distribution is obtained by selecting a number of iterations from the Gibbs sampler for each sum of trees model proportional to its posterior probability, and combining them. The post hoc Gibbs sampler used by BART-BMA is far less computationally expensive than that of BART as it requires only an update for $\mu _{ij}$ and $\sigma$ from the full conditional of each sum of trees model, which is merely a draw from a normal distribution and an inverse-Gamma distribution, respectively (see Sects. D.1, D.2, respectively).

1.1 Appendix D.1: Update of $p(M_j|T_j,R_{\iota ji},\sigma ^2)$

Let $M_j = (\mu _{1j} \ldots \mu _{ij})$ index the $b_j$ terminal node parameters of tree $T_j$, and $R_{kij}$ be the partial residuals for observations k belonging to terminal node i used as the response variable to grow tree $T_j$. BART-BMA assumes that the prior on terminal node parameters is $\mu _{ij}|T_j,\sigma \mathtt {\sim } N(0,\frac{\sigma ^2}{a})$, as in Chipman et al. (1998). The prior distribution of the partial residual is $R_j|\ldots \mathtt {\sim }N(\mu _{ij},\sigma ^2)$.

The full conditional distribution of $M_j$ is then

$$\begin{aligned} p(M_j| T_j,R_{kij},\sigma )&\propto p(R_{kij}|T_j,M_j,\sigma )p(M_j| T_j) \nonumber \\&\propto \prod _{k=1}^{n_i} p(R_{kij}|T_j,M_j,\sigma ) p(M_j |T_j) , \end{aligned}$$

(15)

where k indexes the observations within terminal node i of tree $T_j$ and $n_i$ refers to the number of observations which fall in terminal node i.

The draw from the full conditional of $p(M_j|\ldots )$ is then a draw from the normal distribution

$$\begin{aligned} M_j|T_j,R_{kij},\sigma \mathtt {\sim } N\left( \frac{\sum _{k=1}^{n_i}{R_{\iota ij}}}{n_i+a},\frac{\sigma ^2}{n_i+a}\right) . \end{aligned}$$

(16)

The full conditional of $M_j|\ldots$ depends only on $\sigma$ in the variance parameter, making it slightly more efficient than the update of $M_j$ using the BART prior which depends on $\sigma$ in both the mean and variance parameter.

1.2 Appendix D.2: Update of $p(\sigma ^2)$

BART-BMA performs the update for $p(\sigma )$ in the same way as (Chipman et al. 2010). The full conditional distribution of $\sigma ^2$ is:

$$\begin{aligned} p(\sigma ^2|R_j,T_j,M_j) \propto \prod _{k=1}^n p\left( R_j|T_j,M_j,\sigma ^2 \right) p\left( \sigma ^2\right) ,\nonumber \\ \end{aligned}$$

(17)

where $R_j \mathtt {\sim } N\left( \sum _{j=1}^m g(x_k,T_j,M_j ),\sigma ^2\right)$ and $\frac{1}{\sigma ^2} \mathtt {\sim } \text{ Gamma }(\zeta ,\eta )$, where $\zeta$ and $\eta$ are equal to $\frac{\nu }{2}$ and $\frac{\nu \lambda }{2}$, respectively.

BART-BMA makes the draw for $\sigma ^2$ in terms of the precision $\sigma ^{-2}=\frac{1}{\sigma ^2}$ where $p(\sigma ^{-2}|R_j,T_j,M_j )$ is calculated as:

$$\begin{aligned} \sigma ^{-2}|R_j,T_j,M_j&\mathtt {\sim } \text{ Gamma } \left( \zeta +\frac{1}{2} , \frac{P}{2} + \frac{1}{\eta }\right) , \end{aligned}$$

(18)

where $P=\sum _{k}\left[ Y_k-\sum _{j}g(x_k,T_j,M_j)\right] ^2$ . The next value of $\sigma ^{-2}$ is then drawn from (18), and the value of $\sigma$ is calculated by getting the reciprocal square root of that value.

Appendix E: Greedy tree growing extra details

1.1 Appendix E.1: The PELT algorithm

Univariate changepoint detection algorithms in general search for distributional changes in an ordered series of data. For example, if normality is assumed then such an algorithm may look for changes in the mean or variance of the data. Searching for predictive split points for a single variable in a tree has an equivalent goal, i.e., it is desirable to find split points which maximise the separation of the response variable between the left- and right-hand daughter nodes. For this reason, we use a changepoint detection algorithm called PELT (Pruned Exact Linear Time) in BART-BMA to find predictive split points and greedily grow trees.

PELT was originally proposed to detect changepoints in an ordered series of data $y_{1:n}=(y_1, \ldots , y_n)$ by minimising the function

$$\begin{aligned} \min _\delta \left[ \sum _{\theta =1}^{\Theta +1} \left[ C(y_{(\delta _{\theta -1}+1):\delta _{\theta }})+ D \right] \right] . \end{aligned}$$

(19)

Here, there are $\Theta$ changepoints in the series at positions $\delta _{1:\Theta }=(\delta _1, \ldots , \delta _\Theta )$ which results in $\Theta +1$ segments. Each changepoint position $\delta _{\theta }$ can take the value $1 \ldots n-1$. For example, if a changepoint occurs at position $\delta _1=5$ and another occurs at position $\delta _2=12$, the second segment where $\theta =2$ will contain the values for $y_{(6:12)}$. The function $C(\cdot )$ is a cost function of each segment $\theta$ containing observations $y_{(\delta _{\theta -1}+1):\delta _{\theta }}$. In the results which follow, the cost function used is twice the negative log likelihood assuming that y has a univariate normal distribution. Finally, D is a penalty for adding additional changepoints, default values for which are discussed below.

Table 8 Coverage for out-of-sample 50% prediction intervals and average interval width for BART-BMA, RF using conformal intervals bartMachine and dbarts for the Friedman example

Table 9 Coverage for out-of-sample 75% prediction intervals and average interval width for BART-BMA, RF using conformal prediction bartMachine and dbarts for the Friedman example

PELT extends the optimal partitioning method of Yao (1984) by eliminating any changepoints which cannot be optimal. This is achieved by observing that if there exists a candidate changepoint s where $\delta<s<S$ which reduces the overall cost of the sequence, then the changepoint at $\delta$ can never be optimal and so is removed from consideration (Killick et al. 2012). An algorithm describing how we use PELT to greedily grow trees is described in Algorithm 3:

One disadvantage of using PELT for large datasets is that the number of changepoint detected by the PELT algorithm is linearly related to the number of observations, which can reduce the speed of the BART-BMA algorithm for large n. Our experience is that $D=10\log (n)$ performs well as a general default for the PELT penalty when $n<200$. For larger values of n, we recommend using a higher value for D or the grid search option instead (see Sect. 3.4.2) in order to limit the number of split points detected per variable. We implement a version of PELT which is equivalent to the PELT.meanvar.norm function from the changepoint package in R (Killick et al. 2014). This function searches for changes in the mean and variance of variables which are assumed to be normally distributed. Additional changepoint are accepted if there is support for their inclusion according to the log likelihood ratio statistic.

1.2 Appendix E.2: Updating splitting rules

By default, we choose the best numcp% of the total splitting rules before the tree is grown and only trees using the most predictive splitting rules are considered for inclusion. However, the best splitting rules can also be updated for each internal node i in each tree $T_j$, similarly to how RF creates trees. We have found that updating splitting rules at each internal node generally results in fewer trees $T_j$ being included in each sum of trees model; however, each tree $T_j$ within the sum of trees models averaged over in the final model tends to be deeper and to choose splits that are similar to the primary splits of trees in the RF. We have found that updating the splitting rules at each internal node can in some cases increase the predictive accuracy, but generally at the expense of computational speed.

Appendix F: Out of sample prediction intervals

This appendix shows the results for the calibration of the Friedman example using 50 and 75% prediction intervals (Tables 8, 9).

Fig. 4

Example of experiments to guide default value chosen to determine the size of Occam’s window $o=1000$

Fig. 5

Example of experiments conducted to guide default value chosen to determine the size of the PELT parameter pen where $D=pen\log (n)$ given Occam’s window $o=1000$

Fig. 6

Example of experiments conducted to guide default value chosen to determine the numcp

Appendix G: Choice of default values for BART-BMA

This section will show some of the preliminary investigations which guided the choice of default settings for the BART-BMA algorithm such as the choice of the size of Occam’s window, the penalty on the PELT parameter and the size of sum of tree models to be averaged over. In the results that follow, the following datasets are shown: Ozone, Compactiv, Ankara and Baseball. These were also used as the benchmark datasets for the bartMachine package Kapelner and Bleich (2014b).

For each of the four datasets shown, varying amounts of random noise variables were appended to test the sensitivity of the parameters to the dimensionality of the dataset. In all, 17 different values for the number of random noise variables appended were tested ranging from 100 to 15, 000 so each parameter value of interest was run/tested a total of 68 times.

For the value of Occam’s window, 20 values were evaluated ranging from 100 to 100, 000. A contour plot showing the relative RMSE for the Baseball, Ankara, Compactiv and Ozone datasets can be seen in Fig. 4. Here, the RMSE value for each dataset has been divided by its minimum value which allows for fair comparison across datasets.

In general, we recommend a default value of $OW=1000$ as it seems to work well on the majority of datasets tested as can be seen here (and in other datasets not shown). It was decided that the additional computational complexity involved in setting $OW=10{,}000$ was not worth the marginal gain in accuracy for the datasets tested.

Figure 5 shows the same experiments conducted by ranging the multiple pen used in the PELT penalty $D=pen\log (n)$ from 1 to 20. Here, we can see that a value of $D=10\log (n)$ works well in the majority of the datasets shown.

Figure 6 shows the same datasets where Occam’s window is fixed at its default value of OW $=$ 1000 and the PELT penalty parameter is fixed at $D=10 \hbox {Log}(n)$. Here, 7 increments for numcp were chosen ranging from 5 to 100%.