Soft decision trees for survival analysis

TL;DR

This paper introduces a novel soft survival tree (SST) model that uses soft splitting rules and nonlinear optimization to improve survival analysis, combining flexibility, interpretability, and better performance over existing methods.

Contribution

The paper proposes a new SST model with soft splits and a flexible training formulation, advancing globally optimized survival trees with improved accuracy and interpretability.

Findings

SSTs outperform benchmark survival trees on multiple datasets.

SSTs can incorporate various survival functions, including parametric and spline-based.

The model can be extended to consider group fairness.

Abstract

Decision trees are popular in survival analysis for their interpretability and ability to model complex relationships. Survival trees, which predict the timing of singular events using censored historical data, are typically built through heuristic approaches. Recently, there has been growing interest in globally optimized trees, where the overall tree is trained by minimizing the error function over all its parameters. We propose a new soft survival tree model (SST), with a soft splitting rule at each branch node, trained via a nonlinear optimization formulation amenable to decomposition. Since SSTs provide for every input vector a specific survival function associated to a single leaf node, they satisfy the conditional computation property and inherit the related benefits. SST and the training formulation combine flexibility with interpretability: any smooth survival function (parametric, semiparametric, or nonparametric) estimated through maximum likelihood can be used, and each leaf node of an SST yields a cluster of distinct survival functions which are associated to the data points routed to it. Numerical experiments on 15 well-known datasets show that SSTs, with parametric and spline-based semiparametric survival functions, trained using an adaptation of the node-based decomposition algorithm proposed by Consolo et al. (2024) for soft regression trees, outperform three benchmark survival trees in terms of four widely-used discrimination and calibration measures. SSTs can also be extended to consider group fairness.