Distributed Soft Bayesian Additive Regression Trees

For a p-dimensional vector of predictors $X_{i}$ and a response $Y_{i}(1\leq i\leq n)$ ,the BART model posits

$\displaystyle Y_{i}=f(X_{i})+\varepsilon_{i},\varepsilon_{i}\sim N\left(0,\sigma^{2}\right),i=1,\cdots,n$

(3)

To estimate $f(X)$, a sum of regression trees is specified as

$\displaystyle\widehat{f}(X_{i})=\sum_{j=1}^{m}g\left(X_{i};T_{j},M_{j}\right)$

(4)

$T_{j}$ is the $j^{th}$ binary tree structure and $M_{j}=\left\{\mu_{1j},\ldots,\mu_{b_{j}}\right\}$is the terminal node parameters associated with $T_{j}$ .$T_{j}$ contains information of which bivariate to split on ,the cutoff value ,as well as the internal node’s location. $m$ denote the number of trees which is usually fixed at 200 or 50.

The prior distribution for BART model is $P\left(T_{1},M_{1},\ldots,T_{m},M_{m},\sigma\right)$. Here we assume that $\left\{\left(T_{1},M_{1}\right),\ldots,\left(T_{m},M_{m}\right)\right\}$ are independent with $\sigma$ ,and the $i^{th}$ tree $\left(T_{i},M_{i}\right)$ is independent with the $j^{th}$ tree $\left(T_{j},M_{j}\right)$ when $i\neq j$,so we have

$\displaystyle\begin{aligned} P\left(T_{1},M_{1},\ldots,T_{m},M_{m},\sigma\right)&=P\left(T_{1},M_{1},\ldots,T_{m},M_{m}\right)P(\sigma)\\ &=\left[\prod_{j}^{m}P\left(T_{j},M_{j}\right)\right]P(\sigma)\\ &=\left[\prod_{j}^{m}P\left(M_{j}\mid T_{j}\right)P\left(T_{j}\right)\right]P(\sigma)\\ &=\left[\prod_{j}^{m}\left\{\prod_{k}^{b_{j}}P\left(\mu_{kj}\mid T_{j}\right)\right\}P\left(T_{j}\right)\right]P(\sigma)\end{aligned}$

(5)

So we need to specify the prior for $\mu_{kj}\mid T_{j},\sigma,$ and $T_{j}$.
For the convenience of computation, we use the conjugate normal distribution $N\left(\mu_{\mu},\sigma_{\mu}^{2}\right)$ as the prior for $\mu_{ij}\mid T_{j}$,$(\mu_{\mu}$,$\sigma_{\mu})$can be derived through computation.
The prior for $T_{j}$ is specified by three aspects:

• 1)

  The probability for a node at depth $d$ to split ,given by $\frac{\alpha}{(1+d)^{\beta}}$ .We can confine the depth of each tree by control the splitting probability so that we can avoid overfitting.Usually $\alpha$ is set to 0.95 and $\beta$ is set to 2.

• 2)

  The distribution on the splitting variable assignments at each interior node, default as uniform distribution.Rocková and van der Pas (2017); Linero and Yang (2018) introduced Dirichlet distribution for high dimension variable selection scenario.

• 3)

  The distribution for splitting value assignment,default as uniform distribution.

We use a conjugate prior, here the inverse chi-square distribution for prior of $\sigma$,$\sigma^{2}\sim v\lambda/\chi_{v}^{2}$,the two parameters $\lambda$,$v$ can be roughly derived by calculation.

With the settings of priors $(\ref{equ:s1})$,the posterior distribution can be obtained by

$\displaystyle\begin{aligned} P\left[\left(T_{1},M_{1}\right),\ldots,\left(T_{m},M_{m}\right),\sigma\mid Y\right]\propto&P\left(Y\mid\left(T_{1},M_{1}\right),\ldots,\left(T_{m},M_{m}\right),\sigma\right)\\ &\times P\left(\left(T_{1},M_{1}\right),\ldots,\left(T_{m},M_{m}\right),\sigma\right)\end{aligned}$

(6)

which can be obtained by Gibbs sampling.We need to calculate m successive

$\displaystyle\begin{aligned} P\left[\left(T_{j},M_{j}\right)\mid T_{(j)},M_{(j)},Y,\sigma\right]\end{aligned}$

(7)

where $T_{(j)}$ and $M_{(j)}$ consist of all the trees information and parameters except the $j^{th}$ tree.Then $P\left[\sigma\mid\left(T_{1},M_{1}\right),\ldots,\left(T_{m},M_{m}\right),Y\right]$ can be obtained from sample from inverse gamma distribution with explicit expression.

How to draw from $P\left[\left(T_{j},M_{j}\right)\mid T_{(j)},M_{(j)},Y,\sigma\right]$ ? Note that $T_{j}$, $M_{j}$ depends on $T_{(j)},M_{(j)}$ through $R_{j}=Y-\sum_{w\neq j}g\left(X,T_{w},M_{w}\right)$ ,it’s equivalent to draw posterior from a single tree of

$\displaystyle P\left[\left(T_{j},M_{j}\right)\mid R_{j},\sigma\right]$

(8)

We then split $(\ref{equ:s4})$ in two steps.First we draw from $P\left(T_{j}\mid R_{j},\sigma\right)$,then draw posterior from $P\left(M_{j}\mid T_{j},R_{j},\sigma\right)$.

In the first step,we have

$\displaystyle P\left(T_{j}\mid R_{j},\sigma\right)\propto P\left(T_{j}\right)\int P\left(R_{j}\mid M_{j},T_{j},\sigma\right)P\left(M_{j}\mid T_{j},\sigma\right)dM_{j}$

(9)

,we call $\int P\left(R_{j}\mid M_{j},T_{j},\sigma\right)P\left(M_{j}\mid T_{j},\sigma\right)dM_{j}=P\left(R_{j}\mid T_{j},\sigma\right)$ as marginal likelyhood.Because conjugate Normal prior is employed for $M_{j}$,we can get an explicit expression of the marginal likelihood.We generate a candidate tree $T_{j}^{*}$ from the previous tree structure $T_{j}$ using MH algorithm. We accept the new tree structure $T_{j}^{*}$ with probability

$\displaystyle\alpha\left(T_{j},T_{j}^{*}\right)=\min\left\{1,\frac{q\left(T_{j}^{*},T_{j}\right)}{q\left(T_{j},T_{j}^{*}\right)}\frac{P\left(R_{j}\mid X,T_{j}^{*}\right)}{P\left(R_{j}\mid X,T_{j}\right)}\frac{P\left(T_{j}^{*}\right)}{P\left(T_{j}\right)}\right\}.$

(10)

$q\left(T_{j},T_{j}^{*}\right)$ is the probability for the previous tree $T_{j}$ moves to the new tree $T_{j}^{*}$. The candidate tree $T_{j}^{*}$ is proposed using four type of moves:

• 1)

  Grow,splitting a current leaf into two new leaves.

• 2)

  Prune,collapsing adjacent leaves back into a single leaf.

• 3)

  Swap,swapping the decision rules assigned to two connected interior nodes.

• 4)

  Change,reassigning a decision rule attached to an interior node.

Note that in SBART package,the swap move is not adopted.

Once we have finished sample from $P\left(T_{j}\mid R_{j},\sigma\right)$,we can sample the $k^{th}$ tree the $j^{th}$ leaf parameter $\mu_{kj}$ from a explicit normal distribution.With all the process described above we can iteratively sample from the posterior distribution and obtain valid estimation by merging sampled results.

With the definition of the logistic gating function $\psi(x)$,the probability of going to leaf $\ell$ is

$\displaystyle\phi(x;\mathcal{T},\ell)=\prod_{b\in A(\ell)}\psi(x;\mathcal{T},b)^{1-R_{b}}(1-\psi(x;\mathcal{T},b))^{R_{b}}$

(11)

where $A(\ell)$ is the set of ancestor nodes of leaf $\ell$ and $R_{b}=1$ if the path to $\ell$ goes right at $b$.Here we denote $\phi_{i}$ as the probability vector for the ith sample $x_{i}$ to go to each leaf of the tree.

We get the marginal likelihood with explicit expression

$\displaystyle P\left(R_{j}\mid T_{j},\sigma,\sigma_{\mu}\right)=\frac{|2\pi\Omega|^{1/2}}{\left(2\pi\sigma^{2}\right)^{n/2}\left|2\pi\sigma_{\mu}^{2}\mathrm{I}\right|^{1/2}}\exp\left(-\frac{\|R_{j}\|^{2}}{2\sigma^{2}}+\frac{1}{2}\widehat{\mu}^{\top}\Omega^{-1}\widehat{\mu}\right),$

(12)

where

$\displaystyle\Omega=\left(\frac{\sigma_{\mu}^{2}}{T}\mathrm{I}+\Lambda\right)^{-1},\quad\Lambda=\sum_{i=1}^{n}\phi_{i}\phi_{i}^{\top}/\sigma^{2},\quad\widehat{\mu}=\Omega\sum_{i=1}^{n}R_{i}\phi_{i}/\sigma^{2}$

(13)

So the sufficient statistics for marginal likelihood is $(\sum_{i=1}^{n}\phi_{i}\phi_{i}^{\top},\sum_{i=1}^{n}R_{i}\phi_{i}$) which plays important roles in the distributed system.

Besides the randomized decision rule,SBART use a sparsity-inducing Dirichlet as prior distribution for splitting variables so it can adapted to high dimensional scenario for variable selection.

Distributed computing has attracted increasing attention nowadays along with the real-world datasets become larger and increasingly complex.Some frameworks are developed to handle these situations such as Hadoop,Apache Spark and Message Passing Interface.

Apache Spark is a cluster computing framework that executes the applications much faster than Hadoop. One of it’s advantage is that it can handle situation when some node in the cluster fail. Geirsson (2017) explored a parallel implementation of BART using Apache Spark framework and mentioned that the parallel BART package with MPI support is superior as it minimizes communication costs and the source code is writing in C++ , a faster programming language.

Pratola et al. (2014) built parallel BART with Message Passing Interface support.Their version of BART only allows for growing and pruning steps in the MCMC sampler with a little efficiency lost. They argue that these two steps are sufficient as the trees are small and therefore easily explored.When we use MPI type of distributed computing,the data is partitioned among workers and each worker can work on its own data and communicate with each other by sending messages.We construct the distributed SBART with the help of MPI. Despite the difference generated from random numbers,the distributed version of SBART can get the same result from serial version of SBART.

We can find in the algorithm of SBART that we need to calculate the sufficient statistics 5 times in one tree update.To sample tree structure $\mathcal{T}_{t}$,we need to calculate it twice.To sample bandwidth $\tau_{t}$,we also need twice calculation.And to obtain the tree node parameter estimation $\mathcal{M}_{t}$,another calculation is needed.In fact we can reduce 2 times of sufficient statistics calculation.That is,after we sample tree structure $\mathcal{T}_{t}$,we record the latest result of marginal likelihood and the corresponding tree node parameter estimation $\mathcal{M}_{t1}$ .When we sample $\tau_{t}$,we only need to calculate the sufficient statistics for the new proposed $\tau_{t}$ and the corresponding tree node parameter estimation $\mathcal{M}_{t2}$ .Then use the two group of sufficient statistics to judge whether or not to update the bandwidth.Based on the judgement,we can update the tree with the corresponding parameter estimation $\mathcal{M}_{t1}$ or $\mathcal{M}_{t2}$.

When we sample tree structure $\mathcal{T}_{t}$,the two groups of sufficient statistics are the same for part of the matrix or vector.We should pay attention to part been changed and ignore the unchanged part.By this means,we can accelerate the algorithm and also reduce information size need to transfer.For example,we sample new tree structure $\mathcal{T}_{t}$ with a move of prune.First we get the sufficient statistics $(\sum_{i=1}^{n}\phi_{i}\phi_{i}^{\top},\sum_{i=1}^{n}R_{i}\phi_{i}$) for the old tree structure,we don’t have to calculate all part of the sufficient statistics for the new tree ,we can derive it directly from the old sufficient statistics.In figure $(\ref{FIG1})$ we show an example of prune in which case we can directly obtain the sufficient statistics by merging the information of the two nodes to be pruned together.With change step or grow step,we can only focus on the two nodes involved so we can save some time to calculate and reduce the information need to be transfer between workers and master.

Another try we failed is that we try reduce the process time by saving the $\phi$ for latter use.In big data scenario,$\phi$ become a very large matrix.On the contrary,it need much resources and is slower against our expectation.

We try stop updating the $\tau_{t}$ or attach $\tau_{t}$ with each node and find the estimated accuracy become worse.It’s not beneficial enough to save a little time in exchange of the accuracy.

With all the effort we try to improve SBART in the running time.In Table 1 we show the performance difference between serial SBART and optimized distributed SBART with 2 workers.We can assume that distributed SBART with 2 workers will cost about half of the time optimized SBART with 1 worker need.So we can find that the optimization version is 1.5 times faster for small size up to 2 times faster for large sample size.

Then we outline our distributed algorithm here.Given K workers numbered $0,1,2,\ldots,K-1$,we split the data $(Y,X)$ into $K$ approximately equally-sized portions.The $i^{th}$ data portion $(Y_{i},X_{i})$ resides on worker i .The tree structures $\left(\left(T_{1},M_{1}\right),\left(T_{2},M_{2}\right),\ldots,\left(T_{m},M_{m}\right)\right)$ are kept in every worker.The algorithm follow the master-slave arrangement,where the numbered $0$ core is regarded as a master and the other are slaves.

There is one difference between our algorithm and the work of Pratola et al. (2014).They split the data into $K-1$ portions and leave the master worker with no data.They only assign the MCMC sampler job to the master core.After the MCMC sampler job has been done,the master core will idle all the time and it’s kind of waste computing resources.So we assign one proportion of data to the master core and make full use of the computing resources.It’s helpful when the number of workers is not big.Even when the number of workers is large,we can reduce the data size in the master worker so it can pay more attention to coordinating and communicating with slaves workers.

Considering the draw $(\sigma\mid T_{1},\ldots T_{m},M_{1},\ldots,M_{m},y)$ as a simple example of message passing, we need compute the sufficient statistic $\sum_{i=1}^{n}\epsilon_{i}^{2},$ where $n$ is the total number of observations and $\epsilon_{i}=y_{i}-\sum_{j=1}^{m}g\left(x_{i};T_{j},M_{j}\right)$.So each worker compute the $\epsilon_{i}^{2}$ in its data portion,sum up them and send the results to the master worker.After the master collect all the result and can draw out a new $\sigma$.Then the master will distribute the new $\sigma$ to all the slaves to keep the synchronization of information at different worker.In distributed computing we must know how to decompose sufficient statistics into corresponding job that each worker can carry out with its own data proportion.Not all the statistics need to be processed by this way.For example when sampling the splitting variable $s$,the updated $s$ only need to keep in the master core and has no need to broadcast it out to slave workers.

The main part for the MCMC sampler to update tree structure is illustrate in Figure 2 and we omit the update process for $(s,\sigma,\sigma_{\mu},a)$.The communication path and information size for each step are listed.Compared to the parallel BART its communication cost increase a lot ,that is ,for about $O(d^{2})$ versus $(O(d))$ where $d$ denotes the count of leaves of the tree structure.Fortunately in the BART model we prevent the tree from growing too big.But this increasing communication cost will lead to some problem especially in complicate communication environment.

We then show how we can use distributed SBART to speed up with additional processors. Considering a large dataset $(x,y)$ with 100,000 records.Covariates x is consist of 220 covariables.Response variable is produced by

$\displaystyle y_{i}=10sin(2\pi x_{i1}x_{i2})+(x_{i3}-0.5)^{2}+x_{i4}+x_{i5}+\varepsilon_{i}$

(14)

x is i.i.d. and draws from a uniform distribution $U[0,1]$.$\varepsilon_{i}$ is draw from i.i.d. normal distribution $N(0,1)$.We carry out 1,000 MCMC iterations with the first 500 discarded as burn-in.

Pratola et al. (2014) define the speedup of an algorithm by

$\displaystyle S(n,p)=\frac{T_{seq}}{T_{par}}$

(15)

which is the ratio of the times taken to run two instances of the algorithms. $T_{seq}$ stands for the sequential algorithm’s executing time.$T_{par}$ stands for the parallel algorithm’s executing time. And the efficiency

$\displaystyle E(n,p)=\frac{S(n,p)}{p}$

(16)

is the speedup normalized to the number of cores used in algorithm.n denote the total sample size and p denote the processors involved. Here we can measure the efficiency relative to the speedup of 2 parallel cores.

$\displaystyle E(n,p)=\frac{2*T_{2}}{p*T_{p}}$

(17)

We assign the job to different numbers of cores to see the effect.Table 2 shows the time required to carry out the MCMC draws as a function of the number of processors using this distributed SBART. Here total of 1,000 MCMC iterations were carried out for each run. As expected, the running time decreases with the number of processors.The efficiency is kept in high level which stands for that this algorithm is scalable and we can speed up the algorithm by adding more workers.

If speeding up the algorithm at big data situation is our main reason to choose this distributed algorithm.We can consider another option of embarrassingly parallel.We can easily do this by running copies of the algorithm at different worker. The downside is that the burn-in samples are not parallelized.For example we need to carry out 1000 iterations with the first 500 discarded as burn-in.Now we have 10 workers available.At each workers,we keep a copy of the whole dataset.For each worker we need to carry out 500 burn-in iterations and 50 iterations result independently with no communication with each other with different initial setting.Then combine all the result of the 10 workers together to form 500 valid iterations to poster estimation use.So from this point of view the embarrassingly parallel algorithm is not scalable and no matter how many workers are involved its efficiency is comparable to the distributed SBART with two workers in the previous example.That is one of the reason for us to develop distributed SBART.

Considering dataset $(x,y)$ with 100,000 records.x consist of 10 covariables.These data are produced by

$\displaystyle E(y)=\frac{1}{e^{-(0.4+0.5X_{i1})}}$

(18)

x are i.i.d. draws from a uniform distribution $U[-10,10]$.1,000 MCMC iterations are carried out with the first 500 discarded as burn-in. We generate a test data set with 1,000 records.$X_{1}$ are grid points from -10 to 10 and other x are i.i.d. draws from a $U[-10,10]$ distribution.

In figure 3,we can find that the SBART nearly perfectly match with the true function.And the step function bias of BART can be easily captured which will do harm to the estimation process.From the example it is beneficial to apply SBART model for binary classification problem.

Skin segmentation is a real world data example from UCI data sets(Dua and Graff, 2017). For human object detection, skin segmentation is treated as a pre-processing step followed by the other algorithms. skin dataset is collected by randomly sampling B,G,R values from face images of various age groups (young, middle, and old), race groups (white, black, and Asian) and genders obtained from FERET database and PAL database.Total sample size is 245057 out of which 50859 is the skin samples and 194198 is non-skin samples. We randomly select $80\%$ of the sample as training data set and leave the $20\%$ as test data set.We run the BART and distributed SBART in 2000 iterations with the first 1000 discarded as burn-in.The BART algorithm costs about 845 seconds and the distributed SBART costs about 1896 seconds with 10 workers unless it will cost us more than 5 hours to finish this job with .

We get the confusion matrix as in Table 3 and find SBART perform a little better. That is ,the SBART reduced 7 samples of misclassification in the skin segment in exchange of 10 samples of misclassification in the not skin segment.Because the skin segment ratio is about $21\%$.So we call it a little improvement.

Test set performance for classification problem is measured by area under the Receiver Operating Characteristic curve(AUC), via the ROCR package of Sing et al. (2015). Larger AUC values indicate superior performance, with an AUC of 0.50 corresponding to the expected performance of a method that randomly orders observations by their predictions. A classifier’s AUC value is the probability that it will rank a randomly chosen y = 1 example higher than a randomly chosen y = 0.

Test set performance for this example is measured by AUC.The SBART get the higher AUC result$(99.9934\%)$ than the BART algorithm $(99.9906\%)$.

Another real world data example from UCI data sets(Dua and Graff, 2017) is Higgs data set.The Higgs data has been produced using Monte Carlo simulations. The first 21 features (columns 2-22) are kinematic properties measured by the particle detectors in the accelerator. The last seven features are functions of the first 21 features; these are high-level features derived by physicists to help discriminate between the two classes.We use this data set for two purposes.One is to show its ability to handle big data.Another is to show the advantage of SBART compared to the BART algorithm in classification problem.So we sample 500,000 samples from the data set and apply SBART and BART algorithm to the data set.We randomly sample $90\%$ of the 500,000 as training data and the left $10\%$ as test data.We only analysis the low-level features.We run the BART and distributed SBART in 2000 iterations with the first 1000 discarded as burn-in.

From Figure 4 we can see that the SBART model is superior to the BART model with higher AUC by overcoming the step-wise bias of the BART model.