Efficient gradient boosting for prognostic biomarker discovery

Abstract

Motivation

A gradient boosting decision tree (GBDT) is a powerful ensemble machine-learning method that has the potential to accelerate biomarker discovery from high-dimensional molecular data. Recent algorithmic advances, such as extreme gradient boosting (XGB) and light gradient boosting (LGB), have rendered the GBDT training more efficient, scalable and accurate. However, these modern techniques have not yet been widely adopted in discovering biomarkers for censored survival outcomes, which are key clinical outcomes or endpoints in cancer studies.

Results

In this paper, we present a new R package ‘Xsurv’ as an integrated solution that applies two modern GBDT training frameworks namely, XGB and LGB, for the modeling of right-censored survival outcomes. Based on our simulations, we benchmark the new approaches against traditional methods including the stepwise Cox regression model and the original gradient boosting function implemented in the package ‘gbm’. We also demonstrate the application of Xsurv in analyzing a melanoma methylation dataset. Together, these results suggest that Xsurv is a useful and computationally viable tool for screening a large number of prognostic candidate biomarkers, which may facilitate future translational and clinical research.

Availability and implementation

‘Xsurv’ is freely available as an R package at: https://github.com/topycyao/Xsurv.

Supplementary information

Supplementary data are available at Bioinformatics online.

1 Introduction

Boosting is one of the most influential and successful techniques introduced in the fields of statistics and machine learning for solving complex classification and regression problems. The basic idea is to sequentially learn a combination of multiple weak learners to create a more sophisticated learning model that achieves better prediction performance. AdaBoost, developed in 1995 (Freund and Schapire, 1997), is the first practical implementation in this category. In the early 2000s, a more flexible and effective solution called gradient boosting machine (GBM) was proposed (Friedman, 2001). The GBM generalizes the boosting idea to any differentiable loss functions. By combining the gradient descent algorithm, GBM in each stage identifies an add-on weak learner function by fitting on the negative gradient of the loss function. The main reasons for the widespread application of the GBM today are its flexibility, extensibility and easy implementation. It can be adapted to various loss functions and integrated with different types of weak learners from simple regression to decision trees (i.e. gradient boosting decision tree, GBDT). The GBM has been implemented and integrated into multiple popular R packages, most notably the ‘gbm’ (Greenwell et al., 2007) and ‘caret’ package (Kuhn, 2020). Nonetheless, the GBM faces several drawbacks that might lead to inferior prediction performance in real data analysis. First, like other greedy-search-based algorithms, a GBM and GBDT can converge to local optimal and overfit the training data, especially with a small learning rate and more iterations. Second, a GBM involves more hyperparameters than methods such as random forests (RFs) or elastic net (only two tuning parameters). Therefore, hyper-parameter tuning is more complicated, which often yields less reproducible results in practice.

Extreme gradient boosting (Chen and Guestrin, 2016), (XGBoost or XGB), is an improved version of gradient boosting developed based on the GBDT framework (Hastie et al., 2009; Mason et al., 1999) and has rapidly gained considerable prominence in the field of applied machine learning. Because of its superior prediction performance and computational efficiency, XGB has become a de facto benchmark algorithm in many data science competitions (e.g. Kaggle) and real-world machine-learning projects. XGB significantly mitigates the overfitting issue by introducing extra regularization, built-in tree pruning and subsampling features. As will be discussed in Section 2, the optimization problem in the tree boosting is greatly simplified by the trick in calculating a generalized gain score. In addition, it leverages the multithreading computing capacity offered by modern CPUs. Therefore, XGB scales well to large data scenarios that conventional packages such as gbm cannot accommodate.

LightGBM (Jeschke et al., 2017), hereafter referred to as light gradient boosting (LGB), is another variant of the gradient boosting tree algorithm that has achieved great popularity recently. It integrates multiple novel techniques to optimize computation speed, memory usage and prediction performance jointly. Compared with previous gradient boosting methods, the most notable differences are that LGB adopts a leaf-wise growth strategy (instead of the level-wise strategy) to construct the tree and the gradient-based one-side sampling (GOSS) to find a split. GOSS allows a more efficient (and thus ‘lighter’) implementation of GBDT, which keeps all the data with large gradients and performs random sampling on the data with small gradients. In addition, LGB combines GOSS with a new algorithm called exclusive feature bundling (EFB) to reduce the number of features efficiently. For these reasons, LGB has been recognized as a more suitable alternative to existing gradient boosting implementations in analyzing datasets with extremely high dimensions and considerable sample sizes.

The main scope of this study is to investigate the feasibility of using modern gradient boosting methods, that is XGB and LGB frameworks, for the modeling and prediction of censored survival outcomes. Our work was primarily motivated by a growing demand for more efficient and effective machine-learning methods for prognostic biomarker identification in cancer research. Although XGB and LGB software provide the interface for customizing loss functions, they offer limited support for handling censored survival data. A special form of loss function and prediction evaluation metric needs to be defined for survival outcomes because the outcomes consist of two variables, that is the survival time and censoring information. One crude approximation is to force the survival data to a binary classification problem by dichotomizing the time variable and ignoring the censoring, so that off-the-shelf machine-learning packages can be directly applied. However, it is known that excluding patients with incomplete follow-up will result in an inefficient and biased estimation of coefficients in the survival model. Our study will focus on building the objective function using the Cox proportional hazards (Cox PH) model (Cox, 1972) because it is the most widely used survival model in biomedical and health sciences. A Cox-based objective function is also differentiable and guaranteed to be convex. Of note, despite the newest version of the R package ‘xgboost’ (Chen et al., 2021) also including the survival-based functions (based on the Cox likelihood and the accelerated failure time model), it does not support the selection of different forms of survival loss and evaluation metric functions. Our goal is thus to develop a user-friendly R package (Xurv) that offers a comprehensive range of capabilities relevant to the survival and prognostic biomarker analysis.

The remainder of the paper is organized as follows. In Section 2, we briefly describe the efficient boosting method and provide solutions for Cox’s model and the boosting step for boosting the concordance index. In Section 3, we present simulation results to evaluate the predictive performance of the XGB and LGB survival routines and compare the approaches with the standard Cox regression model and gradient boosting method implemented in gbm. Finally, we apply the developed survival boosting to a melanoma methylation dataset with the goal to identify targeted CpG sites with prognostic values, followed by discussions.

2 Materials and methods

2.1 Gradient boosting overview

Similar to other supervised machine-learning approaches, the ultimate goal of boosting is to find an optimal function of covariates $f^{\mathrm{*}}\left(\mathit{x}\right)$ to predict the outcome Y, by minimizing the loss function $\mathcal{L}(Y,\mathrm{ }f\left(x\right))$. Boosting builds the final predictive model by iteratively combining weak learners that predict the outcome based on a simple model. Thus, it provides an alternative path to build generalized additive models (GAMs): $f\left(\mathit{x}\right)={\sum }_j{f}_j\left(x\right)$, in which each $f_j\left(x\right)$ is a (weighted) weak learner. Boosting performs variable selection implicitly and works properly even with strong multicollinearity and high dimensionality. By incorporating a weighting scheme into resampling steps, boosting is able to focus on more training-informative samples in each step. In the boosted linear regression case, the reweighting is achieved by refitting residuals calculated from previous steps as a surrogate outcome. Gradient boosting generalizes AdaBoost to any smooth loss function $\mathcal{L}(.,\mathcal{ }.)$. In each iteration, gradient boosting trains the base learners by refitting the negative gradients of the loss function. The gradient in this context can be viewed as the pseudo-residuals. In the special case of squared loss, the gradients equal the residuals.

2.1.1 Cox’s survival model for generalized gradient boosting

In biomedical research, the Cox PHs model (Cox, 1972) is the most popular approach for the regression analysis of survival outcomes. In the following, we consider the standard time-to-event (e.g. time to death) data for the ith instance $\left(t_i,{\mathit{x}}_i,{\delta }_i\right),$ where $t_i$ is the observed survival time, ${\mathit{x}}_{\mathit{i}}\mathrm{ }$is the covariate vector and ${\delta }_i\mathrm{ }$is the censoring indicator. Censoring occurs when a time-to-event is not observed during follow-up. We focus on the right censoring, in which the actual event time is no earlier than the observed time. The Cox model defines the hazard function of a subject at time t to be the product of a baseline hazard ${\lambda }_0\left(t\right)\mathrm{ }$ and an exponential function of the covariates,

$$\lambda \left(t|\mathit{x}\right)={\lambda }_0\left(t\right)\exp \left\{H\right({\mathit{x}}_i\left)\right\}.$$ (1)

Here, $H\left(\cdot \right)$ is a risk score function that relates covariates and regression coefficients. In the standard Cox model, the risk score is a linear term, that is H$\left({\mathit{x}}_i\right)={\mathit{x}}_i^T\beta$. In the case of a high-dimensional and non-parametric setting, it is often difficult to specify the functional form of $H(\cdot )$ explicitly. In boosting framework, the target function $H$ is a linear combination of base functions.$\mathrm{ }$For most of the following discussion, we will take derivatives with respect to the risk term $H$, instead of the original parameters $\beta$. A key advantage of the Cox model and its partial likelihood (PL) is that the estimation of $\beta$ does not depend on ${\lambda }_0\left(t\right),\mathrm{ }\mathrm{PL}=\prod {\{\exp (H(\mathit{x}))/{\sum }_{kϵR}\exp (H(\mathit{x}))\}}^{\delta }$. Here, R is the set of the observations at risk at time t. The key to understanding this formula is to recognize that the PL is similar to the conditional probability that a particular study subject is the one that has an event at time t. Here, we want to minimize the negative of the Cox log-PL as the loss function,

$$\mathcal{L}=-\mathrm{logPL}=-{\sum }_i{\delta }_i\left\{H\right({\mathit{x}}_i)-\log [{\sum }_{kϵR}\exp (H\left({\mathit{x}}_{\mathit{i}}\right)\left)\right)\}.$$ (2)

The regression coefficients can be thus estimated by $\widehat{\beta }={\mathrm{argmin}}_{\beta }\left(\mathcal{L}\right(\beta )+R(\beta \left)\right)$, where $R(\beta )$ denotes the regularization terms which constrain coefficients, such as the L1/L2 or group lasso penalty terms (Simon et al., 2013; Yuan and Lin, 2006). Since the Cox objective function is convex, the problem can be efficiently solved by gradient- or subgradient-based algorithms.

Both XGB and LGB start with computing the negative gradient direction of the loss function (working/pseudoresponse or residual) with respect to the current estimate of risk score. The gradient for the Cox loss function is thus,

$$g_i={\delta }_i-\sum _{j=1}^n{\delta }_j{I}_{(t_i\ge t_j)}\frac{\exp \left(H\right({\mathit{x}}_i\left)\right)}{\sum _k^n{I}_{(t_k\ge t_j)}\mathrm{e}\mathrm{xp}\left(H\right({\mathit{x}}_i\left)\right)}.$$

Standard gradient boosting will iteratively optimize the loss function (2) by choosing a weak learner (based on a single or few predictors) that is closest to the negative gradient directions, for example at step m, the optimal basis function ${\eta }^{\left(m\right)}$ can be calculated by

$${\mathrm{argmin}}_{\eta \in H}\sum _{i=1}^n{\left({\eta \left({\mathit{x}}_i\right)-g}_i^{\left(m\right)}\right)}^2.$$

Modern gradient boosting frameworks further consider the second-order Taylor expansion of the loss, which can be viewed as improved variants of Newton boosting$.$ To simplify the expression, we denote ${\tilde{\pi }}_{ij}=I_{(t_i\ge t_j)}\frac{\exp \left(H\right({\mathit{x}}_{\mathit{i}}\left)\right)}{{\sum }_k^n{I}_{(t_k\ge t_j)}\mathrm{e}\mathrm{xp}\left(H\left({\mathit{x}}_{\mathit{k}}\right)\right)}$. $\pi$ can be treated as the estimated absolute probability of failure for subject i at time $t_j$ (given that a failure occurs at $t_j$). The empirical gradient function can now be written as

$$g_i={\delta }_i-\sum _{j=1}^n{\delta }_j{\tilde{\pi }}_{ij}.$$ (3)

As detailed in Supplementary Notes, it can be shown that the second derivative of the Cox PL loss (with respect to H) is

$$s_i=-\sum _{j=1}^n{\delta }_j{\tilde{\pi }}_{ij}(1-{\tilde{\pi }}_{ij}).$$ (4)

Note that both $g_i$ and $s_{\mathrm{i}}$ are empirical terms evaluated at the given data points. At the mth step, we can thus approximate the Cox loss function with functions (3) and (4) through second-order Taylor’s series expansion of the loss around the current function $H^{(m-1)}$,

$${\mathcal{L}}^{\left(m\right)}\approx {\sum }_{i=1}^n\left\{g_i^{\left(m\right)}{\eta }^{\left(m\right)}\left({\mathit{x}}_i\right)+\frac{1}{2}{s}_i^{\left(m\right)}[{{\eta }^{\left(m\right)}\left({\mathit{x}}_i\right)]}^2\right\}+C^{\left(m\right)}.$$ (5)

$C$ only involves the loss term evaluated at $H^{(m-1)},$ which thus can be treated as a constant term at the current iteration. The above equation provides a new optimization objective, which is equivalent to the weighted least squares regression problem. By rearranging terms in Equation (5), the optimal basis function can be expressed as

$${\mathrm{ }\eta }^{\left(m\right)}={\mathrm{argmin}}_{\eta \in H}\sum _{i=1}^n{\left(\frac{1}{2}{s_i^{\left(m\right)}[-\frac{g_i^{\left(m\right)}}{s_i^{\left(m\right)}}-{\eta }^{\left(m\right)}\left({\mathit{x}}_i\right)]}_{\mathrm{ }}^2\right)}^{\mathrm{ }}.$$ (6)

The updated function after step m is then $H^{\left(m\right)}=H^{\left(m-1\right)}+\nu {\eta }^{\left(m\right)}$, where parameter $\nu$ represents the learning rate or step size. The step size moving along the selected direction is determined in a line search step through fitting a linear proportional hazard model (Li and Luan, 2005). The reason that the new formulation is much more efficient than directly optimizing the original loss function in Equation (2) is that most terms are the same for a given iteration and only need to be calculated once, while only the term $\eta \left({\mathit{x}}_i\right)$ needs to be evaluated for each candidate function or new split. The computation is particularly efficient when using a regression tree as the base learner, which will be discussed in the following section.

2.1.2 Efficient tree boosting framework

Tree boosting as proposed originally by Friedman (2001) uses decision trees as base learners. Each leaf (terminal) node in a decision tree is assigned to a prediction value or leaf weight. The tree basis function is defined as ${\eta }^{\left(m\right)}\left(x\right)={\sum }_{l=1}^L{w}_{\mathrm{l}}I(x\in T_l)$, where $w_l$ is the leaf weight and the leaf node indicator function$\mathrm{ }I(x\in T_l)$ defines the structure of a proposed tree. By plugging the tree basis function into Equation (5), the empirical loss function can be rewritten as

$${\mathcal{L}}^{\left(m\right)}=\sum _{l=1}^L\sum _{i\in I_l}^{\mathrm{ }}\left\{g_i^{\left(m\right)}{w}_l+\frac{1}{2}{s}_i^{\left(m\right)}{w_l}^2\right\}=\sum _{l=1}^L\left\{G_i^{\left(m\right)}{w}_l+\frac{1}{2}{S}_i^{\left(m\right)}{w_l}^2\right\},$$

where $G_i$ and $S_i$ are the sum of $g_i$ and $s_i$ in one leaf node. Solving the quadratic function of $w_l$, the optimal solution is thus ${\widehat{w}}_l=-G_i^{\left(m\right)}/S_i^{\left(m\right)}$. The optimal score function is thus $-\frac{1}{2}{\sum }_{l=1}^L\frac{{{[G}_i^{\left(m\right)}]}^2}{S_i^{\left(m\right)}}$. In XGB (Chen and Guestrin, 2016) or LGB (Jeschke et al.), regularization terms are further incorporated into the loss function to control model complexity. For example, if the l2 regularization term of the leaf weights is considered, the optimal score function becomes $-\frac{1}{2}{\sum }_{l=1}^L\frac{{{[G}_i^{\left(m\right)}]}^2}{S_i^{\left(m\right)}+\lambda }+\gamma L$, where $\lambda$ is the l2 regularization term of the leaf weights and $\gamma \mathrm{ }$is the penalty term for the number of terminal nodes. The loss reduction after one split, also known as the ‘gain’ score, is given by

$${\mathcal{L}}_{\mathrm{split}}=\frac{1}{2}\left[\frac{G_L^2}{S_L+\lambda }+\frac{G_R^2}{S_R+\lambda }-\frac{{(G_L+G_R)}^2}{S_L+S_R+\lambda }\right]-\gamma .$$

The l2 regularization term will not only shrink the leaf weight but also, together with the variance term (S), will alter the structure of the final tree. Therefore, the optimization problem is greatly simplified by searching splits that minimize the empirical loss—based on the derived gain function. In the descriptions below, we will use XGB-Cox and LGB-Cox to stand for XGB- and LGB-based algorithms for solving the Cox PL. The current framework can be easily implemented in R and Python and also allows other customized loss functions for survival data as long as it is twice differentiable, such as the smoothed concordance index to be introduced.

2.1.3 Hyperparameter tuning

Similar to XGB and LGB, there are three groups of hyperparameters needed to be determined in Xsurv. The first group is basic parameters defining boosting types (e.g. choice of decision tree or linear model) and parameters for the basic computational environment. The second group is booster parameters such as the number of trees, tree depth, learning rate and regularization terms. Xsurv provides functions to automatically perform cross-validation to identify these hyperparameters. The default learning rate and regularization terms are set from 0.01 to 0.5, with two parameter search strategies (grid search and random search) offered in the main function options. The third group includes more specific learning task parameters that define the objective function and evaluation metrics for validation data.

2.2 Directly boosting the concordance index

One limitation of the Cox model is the assumption of proportion hazards (2). In Xsurv, we also implement an alternative approach based on boosting the C-index (Harrell et al., 1982) directly. The C-index is defined as

$$C=\frac{1}{\left|\mathcal{P}\right|}\sum _{(i,j)\in \mathcal{P}}I(H\left(x_i\right)<H(x_j\left)\right),$$

where $\mathcal{P}$ is the set of orderable pairs and $t_i<t_j$, $H\left(\cdot \right)$ is the assigned risk score. Since the C-index is not differentiable, we use a smoothed concordance index as proposed previously (Chen et al., 2013),

$$SC\mathrm{I}=\frac{1}{\left|\mathcal{P}\right|}\sum _{(i,j)\in \mathcal{P}}\frac{1}{1+\exp (\alpha (H\left(x_i\right)-H\left(x_j\right)\left)\right)}$$

in which $\alpha$ is a hyperparameter that controls the steepness. To implement XGB and LGB, we need the first and second derivative of SCI with respect to $H(\cdot )$. The first derivative of C-index loss function is

$$g_i=\frac{\partial \mathit{SCI}}{\partial H\left(x_i\right)}=\frac{\mathrm{\alpha }}{\left|\mathcal{P}\right|}\left\{\sum _{\left(i,j\right)\in \mathcal{P}}\frac{-\exp \left(\alpha \left(H\left(x_i\right)-H\left(x_j\right)\right)\right)}{{\left[1+\exp \left(\alpha \left(H\left(x_i\right)-H\left(x_j\right)\right)\right)\right]}^2}\right\}+$$

$$\frac{\mathrm{\alpha }}{\left|\mathcal{P}\right|}\left\{\sum _{\left(k,i\right)\in \mathcal{P}}\frac{-\exp \left(\alpha \left(H\left(x_k\right)-H\left(x_i\right)\right)\right)}{{\left[1+\exp \left(\alpha \left(H\left(x_k\right)-H\left(x_i\right)\right)\right)\right]}^2}\right\}.$$

The derivation of second derivative is provided in Supplementary Notes. With these two derivatives, we can optimize the loss function in same way as in the Cox model. In the following sections, we denote the C-index-based boosting methods as XGB-C and LGB-C.

2.3 Survival outcome calibration

One shortcoming of the Cox and C-index-based boosting methods is that the predicted risk scores are only meaningful at the population level. The predicted risk score cannot be immediately interpreted as survival time or probability of individual patients. To facilitate its usage in the personalized medicine setting, Xsurv provides a function to transform predicted scores back to survival time. In addition, we provide an option to output single-patient prognostic group classification (high, medium and low-risk groups), which are more interpretable as illustrated by the example below.

3 Results

In this section, we first present simulation studies that evaluate the performance of the survival gradient boosting methods in terms of risk function estimation and variable selection. We consider survival outcome data generated based on both linear and nonlinear risk score functions, under three different scenarios. We then illustrate the application of Xsurv with an example of discovering prognostic biomarkers in melanoma using methylation data.

3.1 Simulation Scenario 1 (linear model)

In this scenario, we start with a simple proportional hazards model with a linear link function. Let covariates $X=X_1,\dots ,X_p$ be i.i.d. standard normal distributed random variables with dimension $p=100$. The failure time $T$ follows an exponential distribution with mean at $\exp (X_1+X_2+{\dots X}_{10})$. The censoring time C follows an exponential distribution with a mean of $10$, and in order to prevent tail region values from dominating the prediction error, observations exceeding 4 are forced to be censored.

To benchmark the prediction performance of different algorithms, the following six different methods were conducted: Stepwise Cox (Draper and Smith, 1981; Efroymson, 1960; Hocking, 1976), GBM-Cox, XGB-Cox, XGB-C, LGB-Cox and LGB-C. Two metrics are selected to verify our results: the C-index and the integrated brier score (IBS) (Brier, 1950). A total of 100 independent replicates were generated from the same model above and the total sample size n is 1000 in each replicate. The first 800 samples were used as the training data and the remaining 200 as test data. The resulting C-index and IBS are reported in Table  1. As expected, in this simulation setting, stepwise-Cox achieved the best results because it fits the true model. XGB and LGB outperformed GBM with a relatively higher C-index.

Table 1.

Comparisons of predictive performance of the Cox regression, GBM-, XGB- and LGB-based on three simulation scenarios

Model	C-index (SE)	IBS (SE)
Simulation scenario 1
Stepwise-Cox	0.883 (0.009)	0.022 (0.007)
GBM-Cox	0.728 (0.026)	0.054 (0.014)
XGB-Cox	0.782 (0.017)	0.052 (0.010)
XGB-C	0.798 (0.017)	0.045 (0.010)
LGB-Cox	0.770 (0.018)	0.052 (0.010)
LGB-C	0.779 (0.018)	0.049 (0.011)
Simulation scenario 2
Stepwise-Cox	0.587 (0.023)	0.038 (0.014)
GBM-Cox	0.647 (0.026)	0.036 (0.013)
XGB-Cox	0.704 (0.021)	0.033 (0.012)
XGB-C	0.704 (0.023)	0.032 (0.012)
LGB-Cox	0.698 (0.026)	0.032 (0.012)
LGB-C	0.702 (0.024)	0.032 (0.012)
Simulation scenario 3
Stepwise-Cox	0.588 (0.021)	0.039 (0.015)
GBM-Cox	0.638 (0.027)	0.037 (0.014)
XGB-Cox	0.702 (0.020)	0.034 (0.013)
XGB-C	0.702 (0.021)	0.033 (0.013)
LGB-Cox	0.718 (0.020)	0.032 (0.012)
LGB-C	0.708 (0.021)	0.033 (0.013)

3.2 Simulation Scenarios 2 and 3 (nonlinear models)

In the second scenario, we consider nonlinear effects. Let $X=X_1,\dots ,X_p$ be i.i.d. standard normal random variables. We simulated 100 replicates, each with a feature dimension $p=100$ and the sample size of $1000.$ The failure time T follows an exponential distribution with a mean of

$$\exp \left\{2\right[\mathrm{\Phi }\left({[X}_{10}>0.5\right]+X_{20}^2-1+$$

$$\mathrm{\Phi }\left({0.5X}_{30}+X_{40}^2-1\right)+$$

$$\mathrm{\Phi }\left({0.5X}_{50}+X_{60}^2-1\right)+$$

$$\mathrm{\Phi }\left(\sin X_{70}+X_{80}^2-1\mathrm{ }\right)+$$

$$\mathrm{\Phi }\left(\cos X_{90}+X_{100}^2-1\right)\left]\right\},$$

where $\mathrm{\Phi }$ is the standard normal cumulative distribution function. The censoring time $C$ has probability one-third to be 0.02 and probability two-third to be uniform (0, 0.02), and the censoring rate in this case is approximately 30%.

As summarized in Table  1, stepwise-Cox yielded the worst results. XGB and LGB models still showed satisfactory results given the complicated nonlinear generative model. In order to evaluate the algorithms in terms of feature selection accuracy, we examined the top ranked biomarkers based on the feature importance score from all replicates (Fig.  1a). It can be shown that nonlinear signals like $X_{20},X_{40},X_{60},X_{80}$ which are in quadratic forms with large effective sizes had more chances to be detected by XGB and LGB. The quadratic signal $X_{100}\mathrm{ }$is selected with less frequency due to the distortion from $\cos X_{90}.$ Meanwhile, GBM and the stepwise model selected more null or false positive signals if based on the top-ranking feature list.

Fig.1.

Comparisons of top predictive features selected based on 100 replicates in Simulation 2 (a) and Simulation 3 (b). The y-axes list all variables that appeared in top five features in the final predictive model of each training while x-axes represent the frequencies of being selected. The selected ‘True’ signal features are represented by orange bars and the selected ‘Noise’ features are represented by blue ones

In the third simulation scenario, we further considered correlated covariates. This simulation was an extension of Scenario 2 but with covariates $\mathrm{X}={(\mathrm{X}}_1,\dots ,{\mathrm{X}}_{100})$ following a multivariate normal distribution with covariance matrix $\mathrm{V}$, where $V_{ij}={\rho }^{|i-j|}$ and we set $\rho =0.5$ in this scenario. The results are summarized in Table  1 and Figure  1b. Similar to the previous scenario, overall the new gradient boosting models implemented in Xsurv outperformed the baseline models in terms of overall prediction accuracy and feature selection.

3.3 Survival outcome calibration

We further generated a simulation setting to demonstrate the survival calibration function offered by the ‘Xsurv’ package. In this simulation, 100 covariates were generated from a normal distribution with variance 1 and three groups of mean ${\mu }_i$: ${\mu }_1=0,\mathrm{ }{\mu }_2=10$ and ${\mu }_3=15$. We generated 400 samples in each group; hence, the total sample size is 1200. The failure time T follows a Weibull distribution with a shape parameter of 2 and the scale parameter

$$\nu =5\left\{\mathrm{\Phi }\left({[X}_{10}>10\right]+X_{20}^{}-1\right)+$$

$$\mathrm{\Phi }({0.5X}_{30}+{[X}_{40}^{}>5]-1)+$$

$$\mathrm{\Phi }\left({[X}_{50}>10\right]+$$

$${\mathrm{\Phi }[X}_{60}^{}>15\left]\right)+$$

$$\mathrm{\Phi }\left(\left[X_{70}>20\right]+X_{80}^2-1\mathrm{ }\right)+$$

$$\mathrm{\Phi }\left(\right[X_{90}>20]+X_{100}^2-1)\},$$

where $\mathrm{\Phi }$ is the standard normal cumulative distribution function. The censoring time $C$ has probability one-third to be 0.005 and probability two-third to be uniform (0, 0.005). The censoring rate in this case is approximately 30%. For each replicate, we use 1000 samples as training data and the remaining 200 as test data. The risk level is defined based on the value of $\nu$. A larger value of $\mu$ is equivalent to a higher risk of mortality. The real risk level for 1200 samples is defined by its corresponding true $\nu$. For each replicate, samples were divided into three groups based on tertiles of $\nu$, representing ‘Low Risk’, ‘Medium Risk’ and ‘High Risk’, respectively. The result in Figure  2 (based on five replicates) shows that XGB models can successfully classify patients into three risk groups (accuracy around 95%). Similar results were found for LGB models (Supplementary Fig. S1).

Fig. 2.

Survival calibration results from XGB-Cox (left panel) and XGB-C (right panel). The blue dots represent correctly classified (high/medium/low risk group) instances and red dots represent the misclassified instances

3.4 Melanoma methylation dataset

In this section, we demonstrate the use of ‘Xsurv’ through an analysis aiming to identify prognostic methylation biomarkers in melanoma. Matched methylation data (Illumina HumanMethylation 450K array) and gene expression from 470 melanoma samples were downloaded using the R package TCGA2STAT (Liu and Wan, 2015). All methylation values are arcsine transformed on beta values. For illustration purpose, we only select methylation CpG sites that have the strongest correlation (Spearman’s rho >0.3) with the gene expression level of the same gene. This shortlisted CpGs are closely related to the concept of cis-meQTL (Khan, 2018). This step resulted in 9801 candidate CpGs as input data for prognostic biomarker discovery. To minimize patient heterogeneity, we focus the analysis on metastatic melanoma tumor samples. We further filtered 943 CpGs with small coefficients of variations (<0.05). Stage is divided into two groups: stages I, II and I/II NOS into the low stage group and III, IV into the high stage group. The final dataset contains 320 patients with 3 clinical covariates (including sex, age and stage) and 8858 CpGs. The prediction accuracy is evaluated by the C-index and IBS. We compared the results of XGB and LGB with different models such as Lasso-Cox (Tibshirani, 1996, 1997), RF (Ho, 1995) and GBM-Cox. The results are summarized in Table  2. We randomly resampled 100 times from the original dataset with a subsample of 80% of the samples (n = 256 in each experiment). The average C-index and IBS in out-of-bag dataset were calculated. The top ranked CpG biomarkers are shown in Table  2. Among these CpGs, cg13629753 (in gene GBP2) and cg17209284 (in gene SP140L) were the most two selected by different models in Table  2 with 6 and 7, respectively. In addition, they are the only two that were selected by all ‘Xsurv’ models. SP140L is a part of the Speckled Protein (SP) family of chromatin ‘readers’ in humans and mice, and might play a key regulator role in silencing genes that establish immune cell identity and function (Fraschilla and Jeffrey, 2020). Another notable biomarker is cg11429292 (in gene LAG3). LAG3 is a well-known immune checkpoint regulator mostly expressed in tumor infiltrated T cells and a promising target for immunotherapy in melanoma (Pardoll, 2012). We show the top 15 features selected by the LGB-Cox model which have the best C-index using the SHAP value plot (Lundberg and Lee, 2017) in Figure  3a and the aggregated feature selection results for LGB-Cox model in Figure  3b. Results for other models are represented in Supplementary Figure S2. Since there are minor discrepancies for feature ranks between the SHAP- and gain-based results, we also provide the result based on gain in Supplementary Figure S3. We present the recursive partitioning survival tree based on the two most prognostic CpGs (cg13629753 and cg17209284) and stage in Figure  4. Through searching the Human Protein Atlas portal, we found that many top prognostic CpGs identified were in genes that are more pronounced in specific immune cells (https://www.proteinatlas.org/). Together, these results highlight the immunogenic characteristics of cutaneous melanoma and the underlying roles of these CpGs and genes warrant further investigation.

Table 2.

Top prognostic CpGs and corresponding genes selected by each method in the melanoma dataset

	Models
	Lasso-Cox	RF	GBM-Cox	XGB-Cox	XGB-C	LGB-Cox	LGB-C
Top 15 CpGs	cg06942685 ( ZNF542, ZSCAN5A ) cg13984492 (LPPR1) cg17808901 (WDR48) cg19980593 (KCNJ5) cg08036278 (SYVN1) cg14967066 ( IFITM1 ) cg22473973 (BNIP3) cg05371498 (CIITA) cg18419045 (TNXB) cg24670442 (GBP5) cg14750551 ( PARP14 ) cg07745373 (DTX4) cg07535605 (USP14) cg07896558 (NPTX1) cg08639339 (ELK4, SLC45A3)	cg04293930 (TULP1) cg07235805 (PARD6G) cg01550473 (HSPA6) cg12766106 (TRAF2) cg19502867 ( FAM110B ) cg13629753 ( GBP2 ) cg19789466 ( OAS1 ) cg22168987 ( SLC4A2 ) cg05512157 (DIP2B) cg12548899 ( CSTB ) cg24988036 (NECAB3) cg22110267 (ARRB1) cg13577505 (MAPK11) cg21192979 (SLC7A4) cg19692996 (PSMB1)	cg13629753 ( GBP2 ) cg19789466 ( OAS1 ) cg14967066 ( IFITM1 ) cg03270881 (TMSB10) cg07935568 (MLNR) cg20724257 (HLA-DRA) cg00674365 (ZNF471) cg14750551 ( PARP14 ) cg17209284 ( SP140L ) cg27285720 (GBP4) cg16018204 (TDRP, C8orf42) cg19571715 ( OCA2 ) cg22168987 ( SLC4A2 ) cg12548899 ( CSTB ) cg23539753 (SP100)	cg02334987 ( FBXL16 ) cg17209284 ( SP140L ) cg27209571 ( DNER ) cg21533216 (PYURF) cg09709457 (XXYLT1, C3orf21) cg13629753 (GBP2) cg01062113 (PANX1) cg18485596 (SLC16A10) cg06942685 ( ZNF542, ZSCAN5A ) cg27472937 (COLEC11) cg06840243 (UBE2O) cg06576021 (SNAI2) cg25025992 (DNAJA4) cg12336960 (CCDC50) cg22645201 (C7orf58)	cg14967066 ( IFITM1 ) cg13629753 (GBP2) cg16970828 ( UBA7, UBE1L ) cg11429292 ( LAG3 ) cg22874858 ( PDE11A ) cg17209284 ( SP140L ) cg26741686 (ANKRD37) cg27248148 (GGACT) cg10591652 ( FBLL1 ) cg26847438 ( SP140, SP140L ) cg19571715 ( OCA2 ) cg19571715 (OCA2) cg11954384 ( SYNC ) cg09088834 (NINL) cg04706995 (BAZ2B)	cg17209284 ( SP140L ) cg11325273 (DCLK1) cg26464998 (SNED1) cg18523477 (MIR5095, RBM38) cg18119485 (OSBPL9) cg14072140 (DPYD) cg11429292 ( LAG3 ) cg16546864 (FAM207A, C21orf70) cg01136183 (TM7SF4) cg07644368 (CDO1) cg16970828 ( UBA7, UBE1L ) cg13629753 ( GBP2 ) cg19502867 ( FAM110B ) cg14053030 (SIGLEC15) cg04105250 (GAD1)	cg22874858 ( PDE11A ) cg13629753 ( GBP2 ) cg11429292 ( LAG3 ) cg04521957 (SKAP2) cg12074585 (RBP5) cg16970828 ( UBA7, UBE1L ) cg14967066 ( IFITM1 ) cg27209571 ( DNER ) cg15055577 (ALMS1) cg17209284 (SP140L ) cg12446199 (CD247) cg11954384 ( SYNC ) cg10591652 ( FBLL1 ) cg02334987 ( FBXL16 ) cg00767116 (CARS2)
Clinical covariates	Stage	Stage Age	None	Age	Age	Stage	None
C-index	0.548 (0.060)	0.619 (0.046)	0.620 (0.043)	0.625 (0.047)	0.641 (0.042)	0.642 (0.044)	0.631 (0.041)
IBS	0.153 (0.034)	0.142 (0.033)	0.141 (0.028)	0.141 (0.030)	0.141 (0.030)	0.140 (0.030)	0.141 (0.029)

CpGs that were selected by multiple methods are highlighted.

Fig. 3.

Top prognostic CpGs selected by ‘XSurv’ (LGB-Cox) for predicting patient survival in melanoma. (a) SHAP summary plot of top 15 features from the LGB-Cox implementation. Each sample is represented by one dot at each row. (b) Circular barplot showing the frequency of the appearance of the CpG biomarkers in the top 15 features in each subsampling training (based on 100 replicates), where only those that appeared more than five times are displayed

Fig. 4.

Melanoma patient survival prediction based on prognostic CpG biomarkers. (a) The recursive partitioning survival tree based on Stage and the top two robust prognostic CpGs in Table  2. (b) Kaplan–Meier plots comparing the predicted (based on LGB-Cox) and the observed patient overall survival data. (c) Kaplan–Meier plots comparing the patient subgroups stratified by the predicted risk group

3.5 Benchmarking Xsurv versus xgboost

As discussed before, the ‘xgboost’ package only allows the survival analysis to be conducted with the same loss function and evaluation metric. The ‘Xsurv’ package provides more flexibility in allowing multiple commonly used loss functions (such as the C-index) and evaluation metrics for both XGB and LGB model fitting, as well as functions that facilitate automatic parameter tuning and top-biomarker visualization. We performed a benchmark test using four simulation models. The first two synthetic datasets were generated from models in Section 3.1, with P = 50 and Section 3.2 scenario 2, respectively. We also simulated data with a general model similar to Section 6 in Huang (2020),

$$\text{log}\left(\mathrm{T}\right)=-1+\beta X+{ϵ}_{},$$

where ${ϵ}_{}$ follows a normal distribution with mean 0 and variance 4 in the third setting and a nonnormal distribution $\sqrt{12\mathrm{ }}\mathrm{ }\text{logit}\mathrm{ }\left(u\right)/\pi$, where u follows a Uniform(0,1) distribution. We generated 100 replicates for each setting with a sample size of 1000. $\beta$ is a $100\times 1$ vector with first 10 elements generated from a random Uniform(−2, 2) distribution and the others are 0. $X$ was generated with a multivariate standard normal distribution. We randomly select 800 samples as training data and make the prediction on the remaining 200 samples. Supplementary Tables S1 and S2 compare the prediction results from the two packages based on the Cox, AFT and C-index under the same simulation setups. As expected, these implementations achieved very similar performance in all settings, but Xsurv functions can often produce the best results due to flexibility in specifying more types of survival models and boosting methods.

3.6 Computation time

XGB and LGB are both well-known for their fast training speed and high efficiency in dealing with large datasets. We compared the running time of gbm, xgboost, XGB and LGB implemented in ‘Xsurv’ with a 5-fold cross-validation with simulations. We simulated data based on the model described in Section 3.1 with dimensions P = 1000, 5000, 10 000 and 2000, and under four different sample size of 200, 500, 1000 and 5000. Each simulation was repeated 100 times. All computations were done on a Mac Pro server with a 3.3-GHz 12-Core Intel Xeon processor. Results are presented in Supplementary Figure S5. Overall, Xsurv-LGB achieved the best performance in terms of computational efficiency and Xsurv-XGB is slightly faster than xgboost with large training datasets. The results also indicate that the computing time of Xsurv-LGB and Xsurv-XGB remains reasonable even for a genome-wide biomarker scan (e.g. >10 000 genes) with large sample sizes.

4 Conclusion

In this paper, we proposed a modern gradient boosting-based framework, and an R package called Xsurv, to facilitate prognostic biomarker discovery with high-dimensional features and survival outcomes. Xsurv allows efficient survival analysis based on boosting either Cox PH loss or smoothed C-index, which relaxes the PH assumption in real data analyses. Our results in the simulation study suggest both XGB- and LGB-based approaches achieved considerably high prediction accuracy and robust biomarker selection under multiple scenarios. We recommend using LGB implementation when the data dimension is extremely high, given its comparatively higher computational efficiency over XGB. A distinctive feature of our package is that it allows survival time (subgroup) calibration and decision tree visualization, which will greatly aid in the interpretation of learned models. Finally, we applied the developed methods to a melanoma dataset to identify prognostic methylation biomarkers and found many top predictive biomarkers to be indicative of potential immune regulations. Together, our package provides a complementary and integrated tool for biomarker discovery with high-dimensional molecular, clinical or imaging data. Future efforts will be needed to allow the package to handle competing risks and to explore putative mediation effects.