Local Gaussian Process Extrapolation for BART Models with Applications to Causal Inference

Abstract

Bayesian additive regression trees (BART) is a semi-parametric regression model offering state-of-the-art performance on out-of-sample prediction. Despite this success, standard implementations of BART typically suffer from inaccurate prediction and overly narrow prediction intervals at points outside the range of the training data. This article proposes a novel extrapolation strategy that grafts Gaussian processes to the leaf nodes in BART for predicting points outside the range of the observed data. The new method is compared to standard BART implementations and recent frequentist resampling-based methods for predictive inference. We apply the new approach to a challenging problem from causal inference, wherein for some regions of predictor space, only treated or untreated units are observed (but not both). In simulation studies, the new approach boasts superior performance compared to popular alternatives, such as Jackknife+. Supplementary materials for this article are available online.

KEYWORDS:

• Extrapolation
• Gaussian process
• Predictive interval
• Tree
• XBART
• XBCF

Supplementary Materials

Appendix to Local Gaussian process extrapolation for BART: Additional introduction to baseline method Jackknife+ (Barber et al. 2021) in Appendix A; extensive simulation results for XBART-GP and its baseline in Appendix B. (Appendix.pdf)

R-package for XBART-GP: R-package “XBART” is adapted from the original XBART model (https://github.com/JingyuHe/XBART) (He, Yalov, and Hahn 2019; He and Hahn 2021) and contains code to perform local Gaussian process extrapolation on trained XBART model described in this article. Please read the file README contained in the zip file for installation instructions. (XBART.zip, zip archive)

Python-package for XBART-GP: Python-package “XBART” is adapted from the original XBART model (https://github.com/JingyuHe/XBART) (He, Yalov, and Hahn 2019; He and Hahn 2021) and contains code to perform local Gaussian process extrapolation on trained XBART model described in Python. Please read the file README contained in the zip file for installation instructions. (XBART-python.zip, zip archive)

R-package for XBCF-GP: R-package “XBCF” is adapted from the original XBCF model (https://github.com/socket778/XBCF) (Krantsevich, He, and Hahn 2023) and contains code to perform local Gaussian process extrapolation on trained XBCF model described in the article. Please read the file README contained in the zip file for installation instructions. (XBCF.zip, zip archive)

Python code for XBART-GP simulations: The supplemental files contain code to perform simulation experiments described in Section 3.4. (XBART-GP simulation.zip, zip archive)

R code for XBCF-GP simulations: The supplemental files contain code to perform simulation experiments described in Section 4.3. (XBCF-GP simulation.zip, zip archive)

Disclosure Statement

No potential conflict of interest was reported by the authors.

Additional information

Funding

Jingyu He gratefully acknowledges funding for this project fully supported by the Research Grants Council of the Hong Kong Special Administrative Region, China (Project No. CityU 21504921).