PaperlayerView exact revision on arXiv

Source-linked AI summary

Survival regression with accelerated failure time model in XGBoost

Avinash Barnwal, Hyunsu Cho, Toby Dylan Hocking

arXiv:2006.04920v3cs.LGstat.ML

TL;DR

Most tree-based survival models are limited to right-censored outcomes, while binary classifiers require choosing predetermined event-time windows instead of modeling event duration directly. The paper adapts the AFT model and its loss functions to XGBoost, supporting different label-censoring modes, including interval censoring. AFT in XGBoost shows competitive generalization and runtime on interval- and right-censored data, with superior generalization to linear baselines and faster runtime than the nonlinear baseline MMIT.

  • Problem

    Most tree-based survival models are limited to right-censored outcomes, while binary classifiers require choosing predetermined event-time windows instead of modeling event duration directly.

    Source: §1 · p. 3

  • Method

    The paper adapts the AFT model and its loss functions to XGBoost, supporting different label-censoring modes, including interval censoring.

    Source: §1 · p. 3, §2 · p. 4

  • Results

    AFT in XGBoost shows competitive generalization and runtime on interval- and right-censored data, with superior generalization to linear baselines and faster runtime than the nonlinear baseline MMIT.

    Source: §5 · p. 19

  • Takeaways & Limitations

    The paper explores AFT to analyze model coefficients and perform predictive analysis.

    Source: §1 · p. 2

    Two AFT models with nearly identical validation C-index scores can produce highly different mean survival times, so aggregate metrics may miss this sensitivity.

    Source: §4 · p. 18

Abstract

from arXiv · show

Survival regression is used to estimate the relation between time-to-event and feature variables, and is important in application domains such as medicine, marketing, risk management and sales management. Nonlinear tree based machine learning algorithms as implemented in libraries such as XGBoost, scikit-learn, LightGBM, and CatBoost are often more accurate in practice than linear models. However, existing state-of-the-art implementations of tree-based models have offered limited support for survival regression. In this work, we implement loss functions for learning accelerated failure time (AFT) models in XGBoost, to increase the support for survival modeling for different kinds of label censoring. We demonstrate with real and simulated experiments the effectiveness of AFT in XGBoost with respect to a number of baselines, in two respects: generalization performance and training speed. Furthermore, we take advantage of the support for NVIDIA GPUs in XGBoost to achieve substantial speedup over multi-core CPUs. To our knowledge, our work is the first implementation of AFT that utilizes the processing power of NVIDIA GPUs. Starting from the 1.2.0 release, the XGBoost package natively supports the AFT model. The addition of AFT in XGBoost has had significant impact in the open source community, and a few statistics packages now utilize the XGBoost AFT model.

1 Introduction

The introduction motivates adding accelerated failure time (AFT) survival regression to XGBoost because tree-based survival methods offer limited support for different censoring types, while direct time-to-event prediction is useful beyond fixed-window classification. The paper adapts AFT to XGBoost to capture nonlinear patterns, support all censoring modes, and use GPU acceleration.

1 Introduction

  • Cox-PH time prediction requires a nonparametric baseline-hazard estimate, whose computation needs the full training data and is expensive for big data.The baseline hazard depends only on time, while features have multiplicative effects on the hazard.Source: §1 · p. 2
  • Tree-based models capture complex nonlinear feature patterns, and XGBoost’s programming and distributed-computing integrations make survival analysis easier in big-data settings.XGBoost uses second-order loss derivatives to speed convergence.Source: §1 · pp. 2–3
  • Most existing tree-based survival models are limited to right-censored outcomes; Maximum Margin Interval Trees support interval-censored labels.Prior work also includes boosting for AFT and several other tree-based survival models.Source: §1 · p. 3
  • Binary classifiers require predetermined event-time windows, whereas AFT supports direct time-duration prediction and may fit better when the proportional-hazards assumption fails.Cox-PH does not yield a usable time prediction easily, while AFT predictions use fitted parameters and a feature vector.Source: §1 · pp. 2–3
  • The paper adapts AFT to XGBoost, supporting all label-censoring modes, including interval censoring, and evaluating generalization and GPU-accelerated training.Experiments use real and simulated data sets.Source: §1 · p. 3

2 AFT in XGBoost

XGBoost’s AFT formulation replaces linear prediction with a tree-ensemble output and derives censoring-aware losses, derivatives, and numerical safeguards.

2 AFT in XGBoost

  • For gradient boosting, the AFT model uses a tree-ensemble output T(x) in place of the linear feature score.The model relates survival time to input features and a random variable from a known distribution.Source: §2 · p. 4

2.1 Derivation of AFT loss function

  • Training loss uses the density for uncensored labels and the CDF difference across label bounds for censored labels.Summed losses equal the negative log likelihood; an infinite upper bound represents right censoring.Source: §2.1 · pp. 4–5
  • The AFT loss uses the known distribution’s PDF or CDF through a link measuring distance from predicted to observed log survival time.The link is s(y, T(x)) = (ln y − T(x))/σ.Source: §2.1 · pp. 5–6
  • For uncensored labels, normal loss has quadratic tails, logistic loss has symmetric linear tails, and extreme loss has asymmetric tails.For censored labels with finite bounds, the normal and logistic minima lie at the midpoint, while the extreme minimum lies farther along.Source: Figure 1 · p. 6

2.2 Gradient and hessian of the AFT loss

  • XGBoost’s second-order boosting uses the AFT loss’s gradient and Hessian with respect to the tree-ensemble output.The derivatives use the selected distribution’s PDF and CDF, with s(y) = (ln y − u)/σ.Source: §2.2 · pp. 6–7

2.3 Regularization for the AFT loss, to avoid numerical instabilities

  • When predictions are far from the true log-time interval, near-zero CDF differences and large exponentials can produce NaNs or infinities.The small CDF difference affects both the censored loss’s logarithm and derivative denominators.Source: §2.3 · pp. 7–8

2.3.1 Regularization for the loss function (3’)

  • For the loss, replace a CDF difference below 10^-12 with 10^-12.Source: §2.3.1 · p. 8

2.3.2 Regularization for the gradient (5) and the hessian (6)

  • For unstable gradients and Hessians, use specified limiting values, clip gradients to ±15, and assign every point a Hessian of 10^-16.The positive Hessian keeps data points from being ignored by XGBoost.Source: §2.3.2 · p. 8, Table 3 · p. 9

3 Experiments

The experiments assess XGBoost AFT test performance and training speed across censoring settings, alongside hyperparameter search, cross-validation handling of duplicated data, and a CPU–GPU comparison on synthetic data.

3.1 Effectiveness of AFT for interval-censored data

  • For interval-censored labels, the study measured the fraction of predictions within the true log-survival bounds and used nested cross-validation to assess generalization and hyperparameter search.The comparison included survreg, penaltyLearning, and MMIT.Source: §3.1 · pp. 9–10

3.2 Effectiveness of AFT on right-censored data

  • For right-censored data, the study evaluated Uno’s C-statistics with τ set to the 80th percentile of observed survival time and compared XGBoost AFT against two Cox-PH baselines.Nested cross-validation assessed generalization and hyperparameter search.Source: §3.2 · pp. 12, 14

3.1.1 Interval-censored data from supervised changepoint detection problems

  • On ChIP-seq peak-detection data, XGBoost matched SurvReg, penaltyLearning, and MMIT in test accuracy, with runtime near the linear models’ and significantly below MMIT’s.The experiment cast supervised changepoint detection as interval-censored regression.Source: §3.1.1 · p. 11, Figure 2 · p. 13

3.1.2 Synthetic interval-censored data generated from known distributions

  • In Figure 2b, XGBoost and mmit outperformed SurvReg and penaltyLearning in test accuracy; on simulated.model.3, XGBoost achieved 58.5% mean test accuracy versus mmit’s 18% and ran up to 3x faster.The nonlinear models’ additional run time was compensated by higher test accuracy.Source: §3.1.2 · p. 12

3.2.1 Synthetic data with a mix of uncensored and right-censored labels

  • With 20% censoring in Cox-PH-generated data, Scikit-Survival’s Cox-PH had slightly higher C-statistics than the XGBoost models; at 50% and 80%, their scores were similar.For AFT-generated data, XGBoost-AFT with the normal distribution led at 20% censoring, there was no clear winner at 50%, and Scikit-Survival’s Cox-PH led at 80%; XGBoost-AFT ran 2–3× faster in all settings.Source: §3.2.1 · pp. 14–15, Figure 3 · p. 16

3.3 Effect of hyperparameters on model performance

  • Across all runs, random search with 1,000 trials achieved the highest validation accuracy, yet grid search with one or two hyperparameters and fewer than 100 trials yielded higher test accuracy; trying more than 100 combinations made little difference.Source: §3.3 · p. 17

3.4 Efficient model fitting with NVIDIA GPUs

  • For the data simulated.model.3 from Section 3.1.2, five XGBoost models were fit using five cross-validation folds; GPU fitting was 6.1–6.7× faster than CPU fitting in all folds.Source: §3.4 · p. 17
  • Duplicated rows retained their originals’ fold assignments, preserving the fraction of data points in each cross-validation fold.Source: §3.4 · p. 17
  • Table 6 compares CPU and GPU performance using 20 million synthetic data points.The supplied caption gives the comparison's scope but no performance values.Source: Table 6 · p. 18

4 Limitations

The authors identify hyperparameter sensitivity and limitations in survival-curve and confidence-interval support. They describe model stacking as a follow-up approach that addresses these shortcomings.

4 Limitations

  • Two AFT models can have nearly identical C-index values yet produce highly different mean survival times, so aggregate metrics may miss prediction differences.The example varies aft_loss_distribution_scale while keeping other hyperparameters the same.Source: §4 · p. 18
  • XGBoost's predict method returns a point estimate of mean survival time, while the package lacks built-in survival-curve and confidence-interval computation.Survival curves can be computed from the chosen AFT distribution and scale parameter, but may not be well-calibrated.Source: §4 · p. 18
  • XGBSE stacks a second, calibration-amenable model on XGBoost AFT output and fits a Kaplan-Meier estimator to obtain well-calibrated survival curves.The authors report these curves are not sensitive to hyperparameter choice.Source: §4 · pp. 18–19

5 Conclusion

The paper implements AFT in XGBoost and reports competitive generalization and runtime performance on real and simulated data. The implementation became part of XGBoost and is used by open-source statistical packages.

5 Conclusion

  • The authors implemented the Accelerated Failure Time model in the widely used gradient-boosting library XGBoost.The package's ecosystem includes NVIDIA GPU support; the authors note distributed-framework integration as future work.Source: §5 · p. 19
  • AFT in XGBoost shows competitive generalization performance and runtime efficiency for interval-censored and right-censored data, outperforming linear baselines and running faster than mmit.The reported generalization comparison is against survReg and penaltyLearning; the runtime comparison is against the nonlinear baseline mmit.Source: §5 · p. 19
  • After the work became part of XGBoost, packages including XGBSE and MLR3 used its AFT support.The authors identify XGBSE as building on their work and addressing its major shortcomings.Source: §5 · p. 19

A Code

The code section points to the XGBoost package source and the paper's experiment code. The tutorial introduces ranged labels, using a degenerate interval for uncensored observations.

A Code

  • The XGBoost source code and the experiment code are available in separate repositories.The source is in the XGBoost package repository; experiment code is in XGBoostAFTPaperCode.Source: A Code · p. 26

B Tutorial

  • The tutorial represents each label with lower and upper bounds, using [a, a] for uncensored labels.This ranged-label convention is the tutorial's first step.Source: B Tutorial · p. 26

B.1 Using XGBoost AFT in Python

The Python example stores lower and upper label bounds separately and associates them with an XGBoost data matrix. It then configures and trains the AFT model.

B.1 Using XGBoost AFT in Python

  • The Python example stores lower and upper label bounds in separate arrays and passes them to DMatrix as label_lower_bound and label_upper_bound.The example includes uncensored, right-censored, left-censored, and interval-censored labels.Source: §B.1 · p. 26
  • The example configures the objective as survival:aft, selects the aft-nloglik evaluation metric, and sets a normal loss distribution with scale 1.20.It also uses the hist tree method, learning rate 0.05, and maximum depth 2.Source: §B.1 · pp. 26–27
  • The example trains the configured booster for five rounds and evaluates it on the training data.The training call uses dtrain and names the evaluation set train.Source: §B.1 · p. 27

B.2 Using XGBoost AFT in R

The R example represents censored outcomes with lower- and upper-bound label arrays attached to a data matrix. XGBoost's AFT settings select the loss, evaluation metric, and distribution family.

B.2 Using XGBoost AFT in R

  • The example encodes uncensored, right-, left-, and interval-censored labels using lower- and upper-bound arrays attached to the data matrix.It uses setinfo to associate the arrays with the matrix.Source: §B.2 · p. 27

B.3 Note on hyperparameters

  • Set objective to survival:aft and eval_metric to aft-nloglik to use the AFT model in XGBoost.The example then trains the model with xgb.train.Source: §B.3 · p. 27, §B.2 · p. 27
  • Choose normal, logistic, or extreme for aft_loss_distribution; aft_loss_distribution_scale sets the scaling factor σ.The distribution parameter corresponds to Z in the AFT model.Source: §B.3 · p. 27

C Full proof for Definition 3, the gradient and hessian of the AFT loss

The section gives the full proof for the AFT loss gradient and Hessian, deriving the uncensored case and extending the argument to censored data.

C Full proof for Definition 3, the gradient and hessian of the AFT loss

D Effect of hyperparameters on model performance

The hyperparameter analysis finds that 1,000 search trials did not significantly outperform 100 on test accuracy across datasets. However, aggregate accuracy does not establish calibration for individual predictions.

D Effect of hyperparameters on model performance

  • The analysis evaluates generalization using held-out test accuracy for models from each hyperparameter combination.The tables report validation and test accuracy for all searched models; Table D.1 averages both over four outer cross-validation folds.Source: D Effect of hyperparameters on model… · p. 29, Table D.1 · p. 30
  • Across all datasets, 1,000 hyperparameter-search trials provided no significant test-accuracy advantage over 100 trials.The authors conclude that 100 rounds yield a close-to-optimal model when assessed in aggregate by test accuracy.Source: D Effect of hyperparameters on model… · p. 30
  • Aggregate optimality does not ensure individual predictions are well calibrated; models with similar test accuracy can differ substantially on some inputs.Source: D Effect of hyperparameters on model… · p. 30

Related papers

Loading 2006.04920v3…