Scalable piecewise smoothing with BART

TL;DR

This paper introduces ridgeBART, a scalable extension of BART that produces smooth, piecewise functions by integrating localized neural networks, with theoretical guarantees and practical success in spatial data modeling.

Contribution

It develops a scalable MCMC sampler for ridgeBART and establishes theoretical convergence rates for estimating complex functions.

Findings

Effective on synthetic data

Accurate spatial probability estimation in basketball

Nearly minimax-optimal convergence rates

Abstract

Although it is an extremely effective, easy-to-use, and increasingly popular tool for nonparametric regression, the Bayesian Additive Regression Trees (BART) model is limited by the fact that it can only produce discontinuous output. Initial attempts to overcome this limitation were based on regression trees that output Gaussian Processes instead of constants. Unfortunately, implementations of these extensions cannot scale to large datasets. We propose ridgeBART, an extension of BART built with trees that output linear combinations of ridge functions (i.e., a composition of an affine transformation of the inputs and non-linearity); that is, we build a Bayesian ensemble of localized neural networks with a single hidden layer. We develop a new MCMC sampler that updates trees in linear time and establish posterior contraction rates for estimating piecewise anisotropic H\"{o}lder functions and nearly minimax-optimal rates for estimating isotropic H\"{o}lder functions. We demonstrate ridgeBART's effectiveness on synthetic data and use it to estimate the probability that a professional basketball player makes a shot from any location on the court in a spatially smooth fashion.