[Paper Review] Parallel Bayesian Additive Regression Trees

Matthew T. Pratola, Hugh Chipman+4 more|arXiv (Cornell University)|Sep 7, 2013

Gaussian Processes and Bayesian Inference23 references4 citations

TL;DR

This paper introduces a parallel implementation of Bayesian Additive Regression Trees (BART) using single program, multiple data (SPMD) parallel computation via the Message Passing Interface (MPI). By distributing tree construction and MCMC sampling across multiple cores, the method achieves near-linear scalability and enables BART to handle massive datasets—up to 100,000 observations or more—overcoming the computational limitations of the original R-based BayesTree package.

ABSTRACT

Bayesian Additive Regression Trees (BART) is a Bayesian approach to flexible non-linear regression which has been shown to be competitive with the best modern predictive methods such as those based on bagging and boosting. BART offers some advantages. For example, the stochastic search Markov Chain Monte Carlo (MCMC) algorithm can provide a more complete search of the model space and variation across MCMC draws can capture the level of uncertainty in the usual Bayesian way. The BART prior is robust in that reasonable results are typically obtained with a default prior specification. However, the publicly available implementation of the BART algorithm in the R package BayesTree is not fast enough to be considered interactive with over a thousand observations, and is unlikely to even run with 50,000 to 100,000 observations. In this paper we show how the BART algorithm may be modified and then computed using single program, multiple data (SPMD) parallel computation implemented using the Message Passing Interface (MPI) library. The approach scales nearly linearly in the number of processor cores, enabling the practitioner to perform statistical inference on massive datasets. Our approach can also handle datasets too massive to fit on any single data repository.

Motivation & Objective

• To address the computational inefficiency of the original BayesTree R package, which struggles with datasets over 1,000 observations and fails on datasets with 50,000–100,000 observations.
• To enable interactive statistical inference with BART on large-scale datasets by leveraging high-performance computing.
• To develop a scalable, distributed version of BART that maintains the method’s Bayesian uncertainty quantification and robust default priors.
• To support datasets too large to fit on a single machine’s memory by using distributed memory parallelism via MPI.

Proposed method

• Adapt the BART algorithm for SPMD parallelization using the Message Passing Interface (MPI) to distribute computation across multiple processor cores.
• Partition the dataset across multiple compute nodes or cores, with each node independently constructing and updating a subset of trees in the BART ensemble.
• Modify the stochastic search MCMC algorithm to allow parallel updates of tree structures and node parameters across partitions while maintaining posterior consistency.
• Use data sharding and inter-process communication to synchronize sufficient statistics and posterior samples across nodes during MCMC iterations.
• Implement load balancing and efficient I/O handling to manage datasets too large to fit in main memory.
• Ensure the parallelized algorithm preserves the original BART’s robustness and uncertainty quantification through consistent prior specifications and MCMC convergence.

Experimental results

Research questions

• RQ1 Can BART be scaled to large datasets—specifically, 50,000 to 100,000 observations—using parallel computing?
• RQ2 Does the parallel implementation maintain the Bayesian uncertainty quantification and robustness of the original BART method?
• RQ3 What is the performance scaling behavior of the parallel BART algorithm as the number of processor cores increases?
• RQ4 Can the parallel BART framework handle datasets that exceed the memory capacity of a single machine?

Key findings

• The parallel BART implementation achieves near-linear scalability with respect to the number of processor cores, significantly improving runtime performance on large datasets.
• The method successfully handles datasets with up to 100,000 observations, which are infeasible for the original BayesTree R package.
• The algorithm supports datasets too large to fit in a single machine’s memory by distributing data across multiple nodes using MPI.
• The parallel implementation preserves the original BART’s ability to capture uncertainty through variation across MCMC draws and robust default priors.