An Infinite BART model

TL;DR

This paper introduces Infinite BART, a Bayesian ensemble model that automatically determines the number of decision trees and allows different data clusters to have distinct regression functions, enhancing flexibility and interpretability.

Contribution

It generalizes the classic BART model by incorporating an Indian Buffet process prior to automatically select trees and enable cluster-specific regression functions.

Findings

Infinite BART outperforms classic BART in simulated and real data.

The model effectively identifies variable importance and causal effects.

It provides flexible modeling of heterogeneous data clusters.

Abstract

Bayesian additive regression trees (BART) are popular Bayesian ensemble models used in regression and classification analysis. Under this modeling framework, the regression function is approximated by an ensemble of decision trees, interpreted as weak learners that capture different features of the data. In this work, we propose a generalization of the BART model that has two main features: first, it automatically selects the number of decision trees using the given data; second, the model allows clusters of observations to have different regression functions since each data point can only use a selection of weak learners, instead of all of them. This model generalization is accomplished by including a binary weight matrix in the conditional distribution of the response variable, which activates only a specific subset of decision trees for each observation. Such a matrix is endowed with an Indian Buffet process prior, and sampled within the MCMC sampler, together with the other BART parameters. We then compare the Infinite BART model with the classic one on simulated and real datasets. Specifically, we provide examples illustrating variable importance, partial dependence and causal estimation.