[Paper Review] XBART: Accelerated Bayesian Additive Regression Trees

Jingyu He, Saar Yalov+1 more|arXiv (Cornell University)|Oct 4, 2018

Statistical Methods and Inference13 references22 citations

TL;DR

XBART proposes a stochastic hill-climbing algorithm that accelerates Bayesian additive regression trees (BART) by replacing traditional MCMC with a fast, adaptive search strategy, achieving predictive accuracy comparable to standard BART while being orders of magnitude faster and less memory-intensive. It outperforms XGBoost, random forests, and neural networks in accuracy at scale, especially without hyperparameter tuning.

ABSTRACT

Bayesian additive regression trees (BART) (Chipman et. al., 2010) is a powerful predictive model that often outperforms alternative models at out-of-sample prediction. BART is especially well-suited to settings with unstructured predictor variables and substantial sources of unmeasured variation as is typical in the social, behavioral and health sciences. This paper develops a modified version of BART that is amenable to fast posterior estimation. We present a stochastic hill climbing algorithm that matches the remarkable predictive accuracy of previous BART implementations, but is many times faster and less memory intensive. Simulation studies show that the new method is comparable in computation time and more accurate at function estimation than both random forests and gradient boosting.

Motivation & Objective

• To address the high computational cost of standard BART, which relies on slow MCMC sampling and limits its adoption despite strong predictive performance.
• To develop a fast, scalable alternative to MCMC-based BART that maintains Bayesian uncertainty quantification and robustness to hyperparameter choices.
• To enable BART to be used as a default, high-performance method for nonlinear function estimation in settings with complex, unstructured predictors.
• To demonstrate that a stochastic, greedy search strategy can match or exceed the accuracy of MCMC-based BART and other leading tree-based models like XGBoost and random forests.

Proposed method

• XBART replaces the MCMC posterior simulation in BART with a stochastic hill-climbing algorithm that greedily grows trees by selecting splits that maximize the posterior probability of the data.
• The method uses a novel splitting criterion derived from the prior predictive distribution, which incorporates the current estimate of error standard deviation σ to adaptively regularize tree growth.
• It employs a grow-from-root strategy, building trees from the root node and expanding only promising branches, reducing computational overhead.
• Variable selection and split-point adaptation are handled via a predictor-dependent prior that improves model fit and exploration of the tree space.
• The algorithm uses a stochastic proposal mechanism that allows for better exploration of the posterior distribution than deterministic optimizers.
• It integrates a fast, approximate posterior estimation procedure that avoids the burn-in and convergence issues of MCMC, enabling rapid inference.

Experimental results

Research questions

• RQ1 Can a non-MCMC, stochastic hill-climbing algorithm achieve predictive accuracy comparable to MCMC-based BART while being significantly faster?
• RQ2 Does XBART maintain BART’s robustness to hyperparameter choices and provide reliable uncertainty estimates without extensive tuning?
• RQ3 How does XBART compare in accuracy and speed to XGBoost, random forests, and neural networks across diverse data-generating mechanisms?
• RQ4 What is the impact of adaptive regularization via σ-estimation in the splitting criterion on model performance?
• RQ5 Can the proposed algorithm be extended to more complex BART variants, such as Bayesian causal forests or log-linear BART models?

Key findings

• XBART achieved 31% lower root mean squared error (RMSE) than cross-validated XGBoost on average across simulation settings, while being faster in most cases.
• At n = 250,000, XBART was 3.5 times more accurate than untuned XGBoost and 2.5 times faster than the default XGBoost configuration.
• XBART matched or slightly outperformed standard MCMC-based BART in low-noise settings and significantly outperformed it in high-noise settings, likely due to better exploration of the posterior space.
• Random forests and traditional MCMC BART were prohibitively slow at large sample sizes (n > 10,000), whereas XBART scaled efficiently.
• Neural networks and XGBoost performed better in low-noise settings, but XBART maintained superior accuracy across all noise levels without tuning.
• The method’s default hyperparameters, derived from prior elicitation, yielded excellent performance across diverse data-generating processes, confirming its robustness and suitability as a default model.