Research map / Probabilistic and scientific modeling / Uncertainty and Gaussian processes

Very Fast Bayesian Additive Regression Trees on GPU

NeurIPS 2026 · Probabilistic and scientific modeling › Uncertainty and Gaussian processes · approach: Nonparametric and PAC-Bayes Bayesian learning · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Presents bartz, a branchless, JAX-based GPU implementation of the BART MCMC that runs up to 200x faster than single-core CPU and makes BART competitive in time with XGBoost.

Problem

BART is statistically strong but slow on large datasets (n above roughly 10^4–10^5), limiting its use compared to XGBoost.

Approach

Re-express the BART MCMC as a branchless, fixed-depth, fixed-shape sequence of array operations (heap-indexed trees, cached traverse indices and residuals, GROW/PRUNE-only moves) so it can run efficiently on GPU via JAX.

Key results

• bartz on A100 GPU is up to 200x faster than bartz on a single M1 Pro CPU core at high n (n around 1e6 or more).
• At high n, bartz on A100 is about 1000x faster per iteration than XGBoost on L4 GPU, making BART competitive with XGBoost in running time.
• Memory usage is about n(p+n_tree) bytes, roughly 20x less than dbarts (about 24 n(p+n_tree)).

Surprising findings

• Increasing the number of trees well beyond the usual 200 (up to about 1000) may improve predictive performance at high p, an unexplored direction in the BART literature.

Limitations

• No support for binary outcomes/classification yet.
• Runs on a single CPU core or single GPU; no multi-machine sharding.
• GPU memory limits speedup at high p or n_tree (speedup drops to 20-40x).

Experimental setup

Datasets: synthetic regression data (cosine DGP for speed benchmark) synthetic additive-plus-quadratic DGP for RMSE check

Benchmarks: synthetic speed benchmark (time per iteration vs n) synthetic out-of-sample RMSE comparison

Metrics: time per iteration test set RMSE memory usage

Baselines: dbarts XGBoost BART bartMachine XBART

Builds on: BART XBART JAX

Closest papers by meaning

• Bayesian Additive Distribution Regression (NeurIPS 2026 · similarity 0.87)
• BARK: A Fully Bayesian Tree Kernel for Black-box Optimization (ICML 2025 · similarity 0.84)
• Discretization-free Multicalibration through Loss Minimization over Tree Ensembles (NeurIPS 2025 · similarity 0.82)
• Quasi-Bayes meets Vines (NeurIPS 2024 · similarity 0.82)
• Multistage Defer Trees for Hybrid Interpretability: If at First You Can't Succeed, Tree Again (NeurIPS 2026 · similarity 0.81)
• Recursive PAC-Bayes: A Frequentist Approach to Sequential Prior Updates with No Information Loss (NeurIPS 2024 · similarity 0.81)