Bayesian Survival Tree Ensembles with Submodel Shrinkage

Abstract

We consider Bayesian nonparametric estimation of a survival time subject to right-censoring in the presence of potentially high-dimensional predictors. We argue that several approaches, such as random survival forests and existing Bayesian nonparametric approaches, possess several drawbacks, including: computational difficulties; lack of known theoretical properties; and ineffectiveness at filtering out irrelevant predictors. We propose two models based on the Bayesian additive regression trees (BART) framework. The first, Modulated BART (MBART), is fully-nonparametric and models the failure time as the first occurrence of a non-homogeneous Poisson process. The second, CoxBART, uses a Bayesian implementation of Cox’s partial likelihood. These models are adapted to high-dimensional predictors, have default prior specifications, and require simple modifications of existing BART methods to implement. We show the effectiveness of these methods on simulated and benchmark datasets. We also establish that, for a simplified variant of MBART, the posterior distribution contracts at a near-minimax optimal rate in a high-dimensional sparse asymptotic regime.