PaperlayerView exact revision on arXiv

Source-linked AI summary

bartMachine: Machine Learning with Bayesian Additive Regression Trees

Adam Kapelner, Justin Bleich

arXiv:1312.2171v3stat.MLcs.LG

TL;DR

BART lacked an accessible R implementation that combined future prediction with efficient computation and broad analysis features. bartMachine addresses this with a parallelized, feature-rich R package, reducing model-construction time by 50% when in-sample statistics are omitted.

  • Problem

    Existing R BART software lacked convenient future prediction, while faster parallel code was unavailable as an integrated R package.

  • Method

    bartMachine implements parallelized Bayesian additive regression trees in R with prediction, diagnostics, variable selection, interaction detection, and missing-data support.

  • Results

    50% reduction in model-construction time occurs when in-sample statistics are not computed, with approximately 30% speed-up using four versus one core.

  • Takeaways & Limitations

    bartMachine provides a fast, extensive, user-friendly BART implementation intended to broaden access to BART among statistical data analysts.

  • Takeaways & Limitations

    The implementation handles binary classification but not multiclass outcomes, which were planned for a future release.

Abstract

from arXiv · show

We present a new package in R implementing Bayesian additive regression trees (BART). The package introduces many new features for data analysis using BART such as variable selection, interaction detection, model diagnostic plots, incorporation of missing data and the ability to save trees for future prediction. It is significantly faster than the current R implementation, parallelized, and capable of handling both large sample sizes and high-dimensional data.

1. Introduction

bartMachine is an R package implementing a faster, parallelized Bayesian additive regression trees framework designed for accessible, visualization-rich analysis. It addresses limitations of existing BART software while adding prediction, diagnostics, variable-importance, interaction-detection, and related functionality.

  • Package goals and implementation: bartMachine provides a fast, easy-to-use, visualization-rich BART package for R, implemented in Java and integrated through rJava.Its implementation is described as significantly faster than existing approaches and parallelized across as many cores as desired.
  • Package goals and implementation: bartMachine parallelizes model construction and additional operations, including prediction and variable selection.The introduction states that these computations can be divided across available cores.
  • Package goals and implementation: Unlike BayesTree, bartMachine can save trees in memory and predict on test data through convenience functions.The introduction identifies BayesTree’s lack of a predict function as a major drawback because future forecasts require refitting the model.
  • Expanded features: The package adds credible- and predictive-interval plots, convergence diagnostics, permutation-based variable-importance tests, and interaction detection.These features expand visual inspection and variable-importance exploration for BART analyses.
  • Paper organization: The paper introduces BART and its extensions, presents the package and regression and classification examples, and concludes with implementation details in an appendix.The section roadmap covers Sections 2–6 and Appendix A.

2. Overview of BART

BART is a fully Bayesian sum-of-regression-trees approach for flexible nonparametric function estimation, with priors that regularize tree structures and leaf parameters. Its posterior is sampled by Bayesian backfitting, and the implementation supports binary classification but not yet multiclass outcomes.

  • Model formulation: BART estimates unknown functions by summing regression trees that recursively partition predictor space, enabling flexible modeling of interactions and nonlinearities.Each tree partitions predictor space into hyperrectangles, while sums of trees provide greater flexibility than a single tree.
  • Model formulation: The BART model combines m regression trees, each consisting of a tree structure T and terminal-node parameters M, with response noise.The design matrix has n rows and p predictor columns, and the model includes an n × 1 response vector and noise vector.
  • Bayesian specification: BART uses priors on tree structures, leaf parameters, and error variance to regularize the ensemble and prevent any single tree from dominating the fit.The tree-structure prior favors shallow trees; default hyperparameters include α = 0.95 and β = 2.
  • Bayesian specification: Predictor selection is uniformly distributed with probability 1/p in the original formulation, but the implementation permits user-specified variable-selection probabilities.Variables without valid split values receive probability zero, and split values are sampled uniformly from available values.
  • Posterior computation: Posterior draws use Gibbs sampling with Bayesian backfitting, iteratively fitting each tree to residual responses while holding the other m − 1 trees constant.Tree structures are updated with Metropolis-Hastings steps, while leaf parameters and error variance have conjugate posterior distributions.
  • Extensions and limitations: The implementation handles binary classification and identifies multiclass outcomes as planned future functionality.Earlier work explored binary outcomes, while later research extended BART to multiclass problems.

3. The bartMachine package

bartMachine is an R implementation of Bayesian additive regression trees that adds parallel computation, persistent posterior tree ensembles, native missing-data handling, and variable-selection procedures. It also improves model-creation speed through parallelization, memcache, and related computational shortcuts.

  • Package capabilities: The package preserves the posterior distribution of tree ensembles throughout the R session, enabling future prediction without rerunning the Gibbs sampler.The implementation uses an rJava interface to retain the tree ensembles for subsequent prediction and other tree calls.
  • Parallelization: bartMachine parallelizes model creation, prediction, and other features with customizable cores, using one independent Gibbs chain per core during model creation.With four cores, each chain samples 250 burn-in and 250 post-burn-in draws under the default settings, and post-burn-in samples are aggregated.
  • Speed improvements: Model-creation runtime is approximately linear in n, with four cores providing about a 30% speed-up over one core.Turning off in-sample statistics reduces construction time by 50%; memcache should be disabled only above n = 20,000.
  • Speed improvements: For n ≤ 2,000, memcache provides about a 10% speed improvement and can be disabled when memory is limited with little performance degradation.Memcache caches acceptable split variables and values at each node, trading memory for computation speed.
  • Missing data: bartMachine natively incorporates missing data during model creation and future prediction with test data.The implementation uses a missingness strategy known as Missing Incorporated.
  • Variable selection: The package implements variable selection based on variable inclusion proportions, particularly for problems where few covariates influence the response.A cross-validation procedure selects among Local, Global Max, and Global SE thresholds when sparsity is unknown.

4. bartMachine Package Features for Regression

bartMachine provides a ready-to-use BART workflow for regression, including preprocessing, model fitting, diagnostics, cross-validation, prediction uncertainty, variable and interaction analysis, missing-data handling, informed priors, and serialization. Examples on automobile and simulated data demonstrate these capabilities and report strong predictive performance, improved results from informed priors, and scalable model persistence.

  • Regression setup: The regression example predicts log-transformed automobile prices from 25 features using 201 observations after handling missing data and preprocessing predictors.The automobile dataset contains 15 continuous and 10 nominal features; one nominal predictor was dropped and two nominal predictors were coerced to continuous.
  • Model fitting: Default bartMachine settings use 50 trees, 250 burn-in samples, and 1,000 post-burn-in samples, with missing-data handling disabled by default.The defaults also specify α = 0.95, β = 2, k = 2, q = 0.9, ν = 3, and GROW / PRUNE / CHANGE probabilities of 28% / 28% /44%.
  • Model selection: Five-fold cross-validation selects the model with the lowest out-of-sample RMSE, while increasing m > 50 does not substantially improve performance.The selected cross-validated model reports in-sample L1 = 8.18, L2 = 0.68 and Pseudo-R2 = 0.978, and out-of-sample L1 = 21.05, L2 = 4.40 and Pseudo-R2 = 0.858.
  • Diagnostics and uncertainty: Diagnostic tools assess residual normality, homoskedasticity, sampler convergence, and prediction uncertainty through residual plots, convergence plots, and 95% intervals.The example reports errors appearing normal and homoskedastic, a sufficiently burned-in model, and credible or prediction intervals for fitted responses.
  • Informed priors: Informed splitting priors improve simulated out-of-sample RMSE from 1.661159 to 1.232925 when the five true covariates receive five times the weight of 95 useless covariates.Predictor weights are converted internally to splitting probabilities, and informed priors are presented as useful for variable selection and prediction.
  • Extended package features: Additional features include posterior interaction detection, native missing-data prediction with wider credible intervals, and serialization across R sessions.The interaction example identifies x1 and x2 as the dominant interaction, while serialization takes a few minutes and requires 100MB for n = 5000, p = 1000, m = 100, and 5000 posterior samples.

5. bartMachine Package Features for Classification

For dichotomous responses, bartMachine preserves most regression computing settings while removing the σ2 prior parameters and adding classification-specific thresholding and diagnostics. Its classification tools support error-focused covariate tests, probit-scale partial dependence, and credible intervals.

  • The example analyzes 332 Pima Indians, including 109 diabetes diagnoses, using seven continuous predictors and no missing data.
  • Classification uses the same computing parameters as regression, except q and ν cannot be specified because there is no prior on σ2.
  • Lowering prob_rule_class reduces the “YES” class error from 37.6% to 11%, but increases the error rate for the “NO” class.The classification threshold determines whether a probability estimate receives the positive label, so users should cross-validate it against the problem’s objective function.
  • Covariate tests use total misclassification error rather than Pseudo-R2, and the example indicates that age matters for predicting Diabetes after controlling for other predictors.
  • Classification partial dependence plots use the probit of estimated probability, showing a linear increase in Diabetes probability as glucose increases, with 95% credible intervals available.Prediction intervals from the regression case do not exist for classification.

6. Discussion

The article introduces bartMachine as a fast, extensive, user-friendly, open-source R implementation of Bayesian additive regression trees, with code and paper results available for reuse and replication.

  • Discussion: bartMachine implements Bayesian additive regression trees in R as a fast, extensive, and user-friendly package for a broad range of data analysts.The package aims to increase BART’s visibility among statistical audiences through organized, well-documented open-source code.
  • Discussion: The stable package is available on CRAN, while development code is hosted on GitHub under GPL3 and MIT licenses.Scripts in the bart_package_paper repository folder reproduce the paper’s results, tables, and figures.

A. Sampling to Modify Tree Structure

This section explains Metropolis-Hastings sampling for modifying tree structures through GROW, PRUNE, and CHANGE proposals. It calculates the acceptance ratio from transition, likelihood, and tree-structure components, using log-form expressions for numerical accuracy.

  • Proposal mechanisms: Tree structures are modified by growing child nodes, pruning terminal-node pairs, or changing a node’s splitting rule.These operations define the possible tree alterations in the Metropolis-Hastings sampler.
  • Acceptance rule: A proposed tree is accepted when a standard-uniform draw is less than r.The posterior tree probabilities are handled through Bayes’ Rule rather than calculated directly.
  • Ratio components: The tree-structure probability is independent of σ^2.The Metropolis ratio includes a transition ratio and a tree-structure ratio, alongside the likelihood component.
  • Acceptance ratio: The sampler explicitly calculates the acceptance ratio r for GROW, PRUNE, and CHANGE proposals.The calculations are organized around transition, likelihood, and tree-structure ratios.
  • Numerical implementation: The implementation uses log-form expressions for the proposal calculations to improve numerical accuracy.This applies to the expressions used for GROW, PRUNE, and CHANGE proposals.

A.1. Grow proposal · Transition ratio · Likelihood ratio

The BART grow proposal transitions between trees by splitting a terminal node and reversing the move by pruning it. Likelihood calculations factor across terminal nodes and integrate out terminal-node means to obtain the grow-step likelihood ratio.

  • Grow proposal: A grow proposal replaces one current terminal node with two child nodes by selecting a terminal node, predictor, and split point.The selected node is indexed by η, and available predictors and split values depend on the data reaching that node.
  • Transition ratio: The grow transition probability combines the probability of proposing GROW, selecting the terminal node, and selecting its splitting value.The transition ratio is constructed from these proposal components together with the reverse prune probability.
  • Transition ratio: The number of available predictors can be below p when predictors lack at least two unique values within the selected node.This can occur after earlier splits, including splits on dummy variables, and the available split values are adjusted accordingly.
  • Transition ratio: The reverse transition prunes the selected node, and the full transition ratio combines the grow and prune proposal probabilities.The new tree’s second-generation internal nodes determine the number of eligible pruning choices.
  • Transition ratio: When no variable has at least two unique values available for splitting, the GROW probability is set to zero and the proposal is automatically rejected.This condition prevents an invalid grow move.
  • Likelihood ratio: The tree structure assigns responses to terminal nodes, so the likelihood factors into independent probabilities for the observations in those nodes.Rℓ denotes the data in terminal node ℓ, nℓ its observation count, and n the total sample size.
  • Likelihood ratio: BART integrates out each terminal-node mean µℓ rather than conditioning on it, enabling the Gibbs sampler to avoid varying-dimensional continuous-space jumps.The model places a normal prior on the average value of µ.
  • Likelihood ratio: Because only one terminal node changes during a grow move, the likelihood ratio compares that node with its newly created left and right children.The child-node counts are nℓL and nℓR, and substituting the integrated node likelihood yields the grow-step ratio.

Tree structure ratio

The tree structure prior specifies where splits occur and how rules are assigned. Split probabilities depend on node depth and hyperparameters, while rule probabilities depend on available attributes and split points.

  • Tree structure ratio: The tree structure prior governs where splits occur and the rules assigned to nodes.The prior is discussed for the entire tree, including terminal and internal nodes.
  • Tree structure ratio: The split probability at node η is PSPLIT (η) = α/ (1 + dη)β, controlled by α, β, and node depth dη.dη is the number of parent generations of node η.
  • Tree structure ratio: BART assigns rules by selecting from available attributes and unique split points, yielding PRULE (η) = 1/padj(η) × 1/nj·adj(η).The rule probability uses the transition-ratio notation.
  • Tree structure ratio: For a grown node η, the proposal tree adds child nodes ηL and ηR, and the tree-structure ratio combines their nonsplitting probabilities with PSPLIT (η) and PRULE (η).The child-node depths satisfy dηL = dηR = dη + 1.

A.2. Prune proposal … Tree structure ratio

The prune proposal removes both terminal children of a singly internal node, making its transition, likelihood, and tree-structure ratios approximately the inverses of the corresponding grow-proposal ratios. Prune steps exclude single-root trees, and the reverse transition includes a 1/(b −1) term because pruning removes one terminal node.

  • A.2. Prune proposal: Prune selects a singly internal node whose two children are terminal and removes both children.A prune proposal is defined as the opposite of a grow proposal.
  • A.2. Prune proposal: Prune ratios are approximately the inverses of the corresponding grow-proposal ratios.This inverse relationship applies to the proposal ratios described for pruning.
  • A.2. Prune proposal: Prune steps are not considered for trees consisting of a single root node.The exclusion applies specifically to trees without internal child pairs available for pruning.
  • Transition ratio: The forward transition probability factors into the probability of choosing PRUNE and the probability of selecting η to prune from.The supplied expression gives P (T →T∗) = P (PRUNE) P (selecting η to prune from) = P (PRUNE) 1.
  • Transition ratio: The reverse transition must grow from node η, with a 1/(b −1) term because the pruned tree has one fewer terminal node.The reverse direction is analogous to the grow proposal but accounts for the changed number of terminal nodes.
  • Likelihood ratio: The prune likelihood ratio is simply the inverse of the likelihood ratio for the grow proposal.The passage explicitly characterizes the pruning likelihood ratio as an inverse.
  • Tree structure ratio: The prune tree structure ratio is simply the inverse of the tree structure ratio for the grow proposal.The passage explicitly characterizes the pruning tree structure ratio as an inverse.

A.3. Change proposal … Tree structure ratio

The change proposal modifies a singly internal node by selecting a new predictor and valid split value. Its Metropolis-Hastings calculation combines transition, likelihood, and tree-structure ratios, with the latter two simplifying because only the selected node’s children change.

  • A.3. Change proposal: A change proposal selects a singly internal node and changes its rule using a new available predictor and valid split value.Singly internal nodes have two terminal child nodes.
  • Transition ratio: The transition probability includes the change probability, selected-node probability, and probability of selecting the new split value.The first three terms cancel between numerator and denominator when forming the ratio.
  • Transition ratio: The split-value term differs because predictors offer different numbers of unique available split values.The proposal and original trees use n_j*·adj(η*) and n_j·adj(η), respectively, for their available split-value counts.
  • Likelihood ratio: The proposal tree changes only the selected node’s two terminal children, which apportion unexplained responses differently.Residuals are denoted R1· and R2· for the two child nodes, with asterisks marking proposal-tree responses.
  • Likelihood ratio: The likelihood ratio is formed from the response probabilities of the two original and proposal child nodes.The expression is obtained by substituting Equation 9 four times and using algebra; the provided passage also states a condition involving unchanged child response counts.
  • Tree structure ratio: Because the proposal tree has the same structure, the tree-structure ratio only requires the changed node’s children and rule probabilities.The split probabilities remain unchanged because the child nodes retain the same depths.
  • Tree structure ratio: The tree-structure ratio is n_j·adj(η)/n_j*·adj(η*) and is the inverse of the transition ratio.Consequently, only the likelihood ratio must be computed to determine the Metropolis-Hastings ratio r for the change step.

B. Bakeoff

The bakeoff evaluates nine regression data sets using 10-fold cross-validation averaged over 20 replications, finding bartMachine’s predictive accuracy approximately matches the previous implementation.

  • B. Bakeoff: Nine regression data sets were evaluated using out-of-fold RMSE from 10-fold cross-validation averaged across 20 replications.The results are displayed in Table 3.
  • B. Bakeoff: bartMachine performs approximately the same as the previous implementation with regard to predictive accuracy.The comparison concerns predictive accuracy rather than runtime.
  • B. Bakeoff: Average runtime was reported for each algorithm using complete k-fold estimation, with bartMachine run on 4 cores.The runtime results are presented in Table 4 and measured in seconds.
Loading 1312.2171v3…