xgbse: Improving XGBoost for Survival Analysis | by Guilherme Marmerola | Towards Data Science

“There are two cultures in the use of statistical modeling to reach conclusions from data. One assumes that the data are generated by a given stochastic data model. The other uses algorithmic models and treats the data mechanism as unknown.” — Leo Breiman, Statistical Modeling: The Two Cultures
Survival Analysis is a branch of statistics concerned with modeling the duration of events, most importantly providing tools to deal with (and not simply discard) censored data, partial but useful information about durations. For instance, if an apartment is on the market for 90 days, we don't know how long it will take to actually sell it, but we do know that it took at least 90 days (where 90 days is the time of censoring). Survival analysis is useful for a wide range of applications, such as predictive maintenance, customer churn, credit risk, asset liquidity risk, and many others.
The most common output of survival analysis models is a survival curve, which shows the cumulative probability of survival (event not happening) over time. Other models may output risk, which is a number proportional to how fast the curve decays, or an expected time, obtained by direct optimization or by numerically integrating the curve.
At Loft, we're a big user of survival analysis models, as the journey of finding a new home is rather long in Brazil, with several milestones along the way (visits, offers, deeds, renovation, etc). Particularly, as Loft operates an iBuyer (using its own equity to buy, renovate, and sell apartments), we're interested in modeling the time to sell an apartment, as holding it for too long represents the biggest risk we are subject to. The models that produce these estimates are used to make million-dollar decisions and thus have strict requirements on discrimination performance and calibration.
Therefore, we invested a lot of time exploring survival analysis packages. In our quest we found that most implementations embraced one of two cultures:
Models with sound statistical properties (such as calibrated survival curves), but lacking in expressiveness and computational efficiencyHighly efficient and expressive models, but lacking some statistical properties (producing only risk or uncalibrated expected survival times — XGBoost fits here)
Models with sound statistical properties (such as calibrated survival curves), but lacking in expressiveness and computational efficiency
Highly efficient and expressive models, but lacking some statistical properties (producing only risk or uncalibrated expected survival times — XGBoost fits here)
Given this scenario, we embarked on a journey to unite the two cultures in a single package that would meet our requirements: XGBoost Survival Embeddings (xgbse). And now we’re open sourcing it!
We chose XGBoost as our "base model" for its efficient and easy-to-use implementation, and leveraged some tricks to add nice statistical properties to it. We provide more details about XGBoost's shortcomings and these tricks next.
The XGBoost implementation provides two methods for survival analysis: Cox and Accelerated Failure Time (AFT). When it comes to ordering individuals by risk (discrimination), both show competitive performance (as measured by C-index, the ROC AUC equivalent for survival) while being lightning fast.
However, we can observe shortcomings when it comes to other desirable statistical properties. Specifically, three properties are of concern:
prediction of survival curves rather than point estimates
estimation of confidence intervals
calibrated (unbiased) expected survival times
Let us take the AFT implementation as an example (as Cox only outputs risk and not time or survival). The AFT model outputs a value that should be interpreted as the expected survival time for each sample. To compute this time, it assumes an underlying distribution for times and events, controlled by the aft_loss_distribution and aft_loss_distribution_scale hyperparameters. As they control hard-wired assumptions of the model, we would expect that they could drastically change its output.
To confirm this is true, let us perform an experiment using the METABRIC dataset. We will test different values for aft_loss_distribution_scale while keeping aft_loss_distribution as "normal" and check how this will affect model performance. We test the following values: [0.5, 1.0, 1.5]. The following gist gives a rough idea on how to run this simple experiment but you can find the full code here. The results are very interesting, if not alarming:
aft_loss_distribution_scale: 1.5C-index: 0.645Average predicted survival time: 203 days----aft_loss_distribution_scale: 1.0C-index: 0.648Average predicted survival time: 165 days----aft_loss_distribution_scale: 0.5C-index: 0.646Average predicted survival time: 125 days----
In all three scenarios, we can build models with C-index results close to the state of the art [8]. However, our predictions change drastically, with two models showing an average predicted survival time difference of 78 days. This difference would perhaps be understandable if we were analyzing a single sample. But we're looking across the full validation dataset where the average prediction should be fairly stable regardless of which model we use.
If we plot these results alongside an unbiased survival estimator such as the Kaplan Meier we can check that for each step of 0.5 in aft_loss_distribution_scale we move roughly one decile to the right in the curve. Also, we don't see a full survival curve: XGBoost only outputs point time-to-event predictions (no confidence intervals either).
Given these results, we raise the question: "what model should we trust?". Although we could try to optimize for some error measure such as L1-loss [9], we would like more guarantees for applications that are dependent on stable and calibrated time-to-event estimates. Having only point estimates, no confidence intervals, and no "calibration by design" mines trust and prevents shipping survival analysis models to production.
Although in need of an extension for better statistical properties, XGBoost is undoubtedly a powerhouse. C-index results show that the model has great discrimination performance, being competitive with the state of the art. We just need to adapt how we use it.
Besides being leveraged for prediction tasks, Gradient Boosted Trees (GBTs) can also be used as feature transformers of the input data. Trees in the ensemble perform splits on features that discriminate the target, encoding relevant information in their structure. In particular, the terminal nodes (leaves) at each tree in the ensemble define a very nice feature transformation (embedding) of the input data.
This kind of tree ensemble embedding has very convenient properties:
sparsity and high-dimensionality: trees deal with nonlinearity and cast original features to a sparse, high-dimensional embedding, which helps linear models perform well when trained on it. For instance, a logistic regression trained on the embedding (with features as one-hot encoded leaf indices) can show comparable performance to the actual ensemble, with the added benefit of probability calibration (see [1], [2], and [3])supervision: trees also work as a noise filter, performing splits only through features that have predictive power. Thus, the embedding actually has a lower intrinsic dimension than the input data (see [4]). This mitigates the curse of dimensionality and allows a kNN model trained on the embedding (using hamming distance) to have comparable performance to the actual ensemble. With neighbor-sets at hand, we can use robust estimators such as the Kaplan-Meier to get predictions.
sparsity and high-dimensionality: trees deal with nonlinearity and cast original features to a sparse, high-dimensional embedding, which helps linear models perform well when trained on it. For instance, a logistic regression trained on the embedding (with features as one-hot encoded leaf indices) can show comparable performance to the actual ensemble, with the added benefit of probability calibration (see [1], [2], and [3])
supervision: trees also work as a noise filter, performing splits only through features that have predictive power. Thus, the embedding actually has a lower intrinsic dimension than the input data (see [4]). This mitigates the curse of dimensionality and allows a kNN model trained on the embedding (using hamming distance) to have comparable performance to the actual ensemble. With neighbor-sets at hand, we can use robust estimators such as the Kaplan-Meier to get predictions.
These two properties are the foundation upon which we built xgbse. We show how we put them to good use next.
Our first approach, XGBSEDebiasedBCE, takes inspiration from the multi-task logistic regression method in [5], the BCE (binary cross-entropy) approach in [6], and the probability calibration ideas from [1], [2] and [3].
It consists of training a set of logistic regressions on top of the embedding produced by the underlying XGBoost model, each predicting survival at different user-defined discrete time windows. The classifiers remove individuals as they are censored, with targets that are indicators of surviving at each window.
As the embedding is sparse and high dimensional, the linear models actually can improve the ensemble performance, while also adding nice statistical properties to it, such as probability calibration (as logistic regressions are calibrated by design).
The naive approach tends to give biased survival curves, due to the removal of censored individuals. Thus, we made some adaptations such that logistic regressions estimate the di/ni term (point probabilities) in the Kaplan-Meier formula and then use the KM estimator to get nearly unbiased survival curves.
This way, we can get full survival curves from XGBoost, and confidence intervals with minor adaptations (such as performing some rounds of bootstrap). Training and scoring of logistic regression models is efficient, being performed in parallel through joblib, so the model can scale to hundreds of thousands or millions of samples.
As explained in the previous section, even though the embedding produced by XGBoost is sparse and high dimensional, its intrinsic dimensionality actually should be lower than the input data. This enables us to "convert" XGBoost into a nearest neighbor model, where we use hamming distance to define similar elements as the ones that co-occurred the most at the ensemble terminal nodes. Then, at each neighbor-set we can get survival estimates with robust methods such as the Kaplan-Meier estimator.
We recommend using dart as the booster to prevent any tree to dominate variance in the ensemble and break the leaf co-occurrence similarity logic. We built a high-performing implementation of the KM estimator to calculate several survival curves in a vectorized fashion, including upper and lower confidence intervals based on the Exponential Greenwood formula [7].
However, this method can be very expensive at scales of hundreds of thousands of samples, due to the nearest neighbor search, both on training (construction of search index) and scoring (actual search).
As a simplification to XGBSEKaplanNeighbors, we also provide a single tree implementation. Instead of doing expensive nearest neighbor searches, we fit a single tree via XGBoost and calculate KM curves at each of its leaves.
It is by far the most efficient implementation, able to scale to millions of examples easily. At fit time, the tree is built and all KM curves are pre-calculated, so that at scoring time a simple query will suffice to get the model’s estimates.
However, as we’re fitting a single tree, predictive power may be worse. That could be a sensible tradeoff, but we also provide XGBSEBootstrapEstimator, a bootstrap abstraction where we can fit a forest of XGBSEKaplanTree's to improve accuracy and reduce variance.
Now we return to the first example and check how XGBEmbedKaplanNeighbors performs:
aft_loss_distribution_scale: 1.5C-index: 0.640Average probability of survival at [30, 90, 150] days: [0.9109, 0.6854, 0.528]----aft_loss_distribution_scale: 1.0C-index: 0.644Average probability of survival at [30, 90, 150] days: [0.9111, 0.6889, 0.5333]----aft_loss_distribution_scale: 0.5C-index: 0.650Average probability of survival at [30, 90, 150] days: [0.913, 0.6904, 0.5289]----
As measured by the average probability of survival in 30, 90 and 150 days the model shows good results, with comparable C-index results and predictions being very close and independent of aft_loss_distribution_scale choice. Visually, the comparison of the average model predictions to a Kaplan Meier yields much better results:
Much less bias, full survival curves and stability! Although is too harsh to claim that the problem is solved, we believe that the we could build a strong survival analysis package and bring the "two cultures" much closer together.
XGBoost Survival Embeddings shows great results in several survival analysis benchmarks, outperforming vanilla XGBoost and parametric methods in most cases. It is pip installable and scikit-learn compatible. There are more features to try, such as confidence intervals, extrapolation of survival curves, explainability through prototypes, metrics, and more. We’re committed to improving the lib and adding new content in the future.
For more details, please check the the library on github, the docs, and examples! We invite the community to contribute. Please help by trying it out, reporting bugs, and letting us know what you think!
By Davi Vieira, Guilherme Marmerola, Gabriel Gimenez and Vitor Estima.
[1] X. He, J. Pan, O. Jin, T. Xu, B. Liu, T. Xu, Y. Shi, A. Atallah, R. Herbrich, S. Bowers, and J. Q. Candela. Practical Lessons from Predicting Clicks on Ads at Facebook (2014). In Proceedings of the Eighth International Workshop on Data Mining for Online Advertising (ADKDD’14).
[2] Feature transformations with ensembles of trees. Scikit-learn documentation at https://scikit-learn.org/.
[3] G. Marmerola. Calibration of probabilities for tree-based models. Personal Blog at https://gdmarmerola.github.io/.
[4] G. Marmerola. Supervised dimensionality reduction and clustering at scale with RFs with UMAP. Personal Blog at https://gdmarmerola.github.io/.
[5] C. Yu, R. Greiner, H. Lin, V. Baracos. Learning Patient-Specific Cancer Survival Distributions as a Sequence of Dependent Regressors. Advances in Neural Information Processing Systems 24 (NIPS 2011).
[6] H. Kvamme, Ø. Borgan. The Brier Score under Administrative Censoring: Problems and Solutions. arXiv preprint arXiv:1912.08581.
[7] S. Sawyer. The Greenwood and Exponential Greenwood Confidence Intervals in Survival Analysis. Handout on Washington University in St. Louis website.
[8] H. Kvamme, Ø. Borgan, and I. Scheel. Time-to-event prediction with neural networks and Cox regression. Journal of Machine Learning Research, 20(129):1–30, 2019.
[9] H. Haider, B. Hoehn, S. Davis, R. Greiner. Effective Ways to Build and Evaluate Individual Survival Distributions. Journal of Machine Learning Research 21 (2020) 1–63.