Deep learning-based ranking method for subgroup and predictive biomarker identification in patients

Cohort data set

The analysis population in this study comprised AbbVie PIONEER I (ClinicalTrials.gov numbers: NCT01468207) and II (ClinicalTrials.gov numbers: NCT01468233) patients who are re-randomized to either continuation of adalimumab weekly dosing or withdrawal from adalimumab (placebo) in period B after initial treatment of adalimumab weekly dosing for 12 weeks, Fig. 1. The proportion of female participants is 63.8% in PIONEER I and 67.8% in PIONEER II. The mean (standard deviation) age of participants is 37.0 (11.1) years in PIONEER I and 35.5 (11.1) years in PIONEER II. A total of 199 patients (99 continued adalimumab weekly dosing, 100 withdrawn from adalimumab weekly dosing) in period B on the integrated data from the two studies were included for this analysis. Based on the study team discussion, clinically relevant candidate biomarkers are included: % reduction in AN count at week 12, AN count at week 12, AN count at week 0, draining fistula count at week 0, draining fistula count at week 12, reduction in draining fistula count at week 12, abscess count at week 0, reduction in abscess count at week 12, Hurley stage at week 0, smoking status, abscess count at week 12, initial HiSCR responder status at week 12 and concomitant use of antibiotics.

Ethics approval and informed consent

The two clinical trials included in this study are conducted in accordance with the International Conference on Harmonisation guidelines, applicable regulatory requirements, and the principles of the Declaration of Helsinki. The study protocols (AbbVie protocol number M11-313; EudraCT number 2011-003400-20) are developed collaboratively by the investigators and the sponsor (AbbVie) and are approved by the independent ethics committee or institutional review board at each participating site. Written informed consent is obtained from all participants prior to enrollment, and this consent included permission for secondary analyses of data collected during the clinical trials. While the names of the specific review boards and institutions are not publicly disclosed in full, ethical oversight and site participation details are documented in the original publication of the trials⁴ and on file with the study sponsor (AbbVie).

Subgroup identification framework

We denote the observed data by {(Xi, A_i, Y_i), \(i=1,\ldots ,n\)} consisting of \(n\) independent patients, where \({Y}_{i}\) denotes outcome, X_i and A_i represents covariate of biomarkers and treatment assignment for ith subjects. We adopt the Neyman-Rubin potential outcome framework in causal inference^21,22. In this framework, only one of the potential outcomes can be observed, that is, \({Y}_{i}\) \(={\frac{1}{2}\left(\right.1+A}_{i}\))\({Y}_{i}(1)+\frac{1}{2}\)(\({1-A}_{i}\)) \({Y}_{i}(-1)\), where Y_i (1) and Y_i (-1) are the potential outcomes if the patient i receives a treatment \(({A}_{i}=1)\) and a control \(({A}_{i}=-1)\), respectively. Let (X, A, Y) denote identically distributed copies of observed data, the completely unspecified regression model is formulated as follows:

where \(Z\left(X\right)=\frac{1}{2}\left[E\left(Y|A=1,X\right)-E\left(Y|A=-1,X\right)\right]\) is a contrast function that reflects treatment effects given \(X\) and \(H\left(X\right)=\frac{1}{2}\left[E\left(Y|A=1,X\right)+E\left(Y|A=-1,X\right)\right]\) is a function that reflects the prognostics effect of \(X\). The estimator of \(Z\left(X\right)\) is our interest in subgroup identification, as it reflects treatment effect heterogeneity. Our goal is to estimate the treatment difference Z(X) as the metric for summarizing ITRs without the need to estimate \(H(\cdot ),\) then biomarker importance can be assessed based on their influence on \(Y\) only via \(Z(X)\). Nonetheless, it’s important to note that not all subgroup identification methodologies are designed to target \(Z(X)\); instead, certain methods might focus on deriving a useful transformation of \(Z(X)\)^19,23.

We construct a personalized benefit scoring system, defined as \(f(X)\) with the following two properties: i) \(f(X)\) is monotone in the treatment \(Z\left(X\right)\); ii) it has a threshold value \(c\), such that when f(X) > c, it implies that the treatment is more effective than control. In this work, we consider \(c=0.\) Therefore\(,\) the sign {f(X)} can be used to construct optimal ITRs and predict which of two treatments will have a better outcome. When sign {f (X)} >0\(,\) patients are assigned to treatment group. The importance of each biomarker can be assessed regarding its contribution to the predictive modeling of ITRs. To ensure the identifiability of ITRs, we make following two assumptions by using the standard strong ignorability condition^24,25,26: {\(Y\left(1\right)-\), \(Y\left(-1\right)\)} \(\perp\) \({A|X}\) and 0 < P(A = 1|X) < 1 for all \(x\).

DeepRAB model

We introduce a DNN-based method for estimating ITRs and identifying predictive biomarkers. DeepRAB is a nonlinear model designed to capture \(Z(X)\) using biomarkers values as inputs and disease outcomes as output. DeepRAB consists of three main components. First, it incorporates the A-learning approach¹⁹ into the loss function to optimize ITRs. Second, it features a biomarker selection layer within the encoder layer, which compresses the input into a lower-dimensional representation and selects the biomarkers with the most impact on Y through \(Z(X).\) The third component is a multi-layer perceptron (MLP) known as the hidden layers within decoder layer. These layers model potentially non-linear effects of the covariates.

Let us set the pre-fixed number of \(k\) nodes in the biomarker selection layer, the output of encoder layer, denoted by z⁽¹⁾, can be expressed as z⁽¹⁾ = B^(0),Tx, where z⁽¹⁾ ∈R^k and B⁽⁰⁾ ≡[β_k⁽⁰⁾,…β_k⁽⁰⁾] ∈R^p×k. Here, \({z}_{i}^{(1)}\)=β_i^(0),Tx = x₁ \({\beta }_{i1}^{(0)}+\ldots +{x}_{p}{\beta }_{{ip}}^{(0)}\) for \(i=1,\ldots ,k\) represents the output of the ith node. Notably, this layer is used to select a user-specified set of k biomarkers that are deemed to be the most informative for predicting \(f(X)\). To accomplish this, we adopt the feature selection technique proposed in Balın et al¹⁸. which involves generating a \(p\)-dimensional vector \({\beta }_{i}^{(0)}\) using:

where \({\beta }_{{ij}}^{\left(0\right)}\) corresponds to the \({jth}\) element in \({{\boldsymbol{\beta }}}_{i}^{\left(0\right)}.\) The vectors α and g are p-dimensional and training parameters, and all elements of α are strictly greater than zero, while all elements of g are drawn from a Gumbel distribution²⁷. The temperature parameter \(T\) takes on positive values. In this way, we obtain \({{\mathrm{lim}}}_{T\longrightarrow 0}{{\boldsymbol{\beta }}}_{i}^{\left(0\right)}={[{\mathrm{0,0}},\ldots {\mathrm{1,0,0}},\ldots ,0]}^{{{\rm{T}}}}\) with probability \(P=\frac{{\alpha }_{j}}{{\sum}_{t}{\alpha }_{t}}.\). We sample a \(k\times p\) dimensional matrix \({B}^{0}\) for each of the k nodes in a similar manner. As a result, each node in the biomarker selection layer outputs one selected biomarker, resulting in a total of \(k\) selected biomarkers.

The decoder layers take z⁽¹⁾ as the input and are composed of \(h-1\) hidden layers, where \(h\) is tunning parameters. The expressions for the outputs of each hidden layer, denoted by \({{\boldsymbol{\ d}}}^{(j)},{{\rm{where}}}\) \(j=1,..,h-1,\) can be formulated as follows:

where n_j \({{\rm{and}}}\) ϕ_j denotes the number of nodes and activation function in jth hidden layer, respectively. Then, we can write output f (\(x\)):

where ϕ_h is the activation function which typically a logistic function for binary outcomes and linear function for continues outcomes. The loss function of the model is defined as follows:

where \({{\boldsymbol{\theta}}} \equiv ({W}_{1\times {n}_{h}}^{\left(h\right)},\,\ldots ,\,{W}_{{n}_{1}\times k}^{\left(1\right)},\,{{\boldsymbol{\ b}}}_{1\times 1}^{\left(h\right)},\ldots ,{b}_{{n}_{1}\times 1}^{\left(1\right)},{{\boldsymbol{\alpha}}} ,{{\boldsymbol{g}}})\) represents the training parameters, and π(X)=P(A=1│X) is propensity score. The sign {\(f(\widetilde{\theta },X)\)} is used to construct optimal ITRs. Of note, the function \(M\)(\(u,v\)) varies depending on the outcome Y. In the original work by Chen¹⁹, \(M\left(y,v\right)\) is required to meet the following conditions: 1) \(M\left(y,v\right)\) is convex in \(v\) and 2) \(V\left(y\right):= M\left(y,0\right)\) is monotone in \(y\). These requirements are sufficient for Fisher consistent subgroup identification^11,28. In this work, we select \(M\left(y,v\right)\) as follows: M(u, v)=(u–v)² for continuous outcomes and \(M\)(\(u,v\))\(=u\log \left(1+\exp \left(-v\right)\right)\) for binary outcomes. It can be readily verified that these choices fulfill both the convexity in \(v\) and the monotonicity in \(y\) conditions. A visual representation of the DeepRAB is presented in Fig. 2.

Propensity scores

The propensity scores π(X_i) is nuisance parameter and unknown in observational studies. However, in randomized trials, the propensity scores are often known. The special case is that \(\pi \left({X}_{i}\right)=\frac{1}{2}\),\({{\rm{for}}}\) all i=1,…,n, when the samples are subjected to a 1:1 randomization ratio. In non-randomized trials, we employ logistic regression to the data (\(X,A\)) to estimate \(\pi \left({X}_{i}\right)\).

Simulation framework

We consider three simulation scenarios: one involving linear functions and two involving nonlinear functions, introducing greater complexity to the data. For each scenario, we assess performance using two sample sizes, \(N=1000\) and \(N=400\), with the smaller sample size reflecting conditions commonly encountered in real datasets.

Simulation scenario I

We first consider a linear simulation design to generate data for a continuous outcome using the following model:

where ϵ~N(0,1) represents error terms, and covariate vector \(X=({{X}_{1},{X}_{2},\ldots {X}_{10}})^{{{\rm{T}}}}\) is generated from a multivariate normal distribution with a mean of 0, variance of 2, and pairwise correlations of 0.2 between each biomarker. The treatment assignment variable \(A\) is drawn from a Bernoulli \(\left(0.5\right)\) distribution, with \(A=0\) indicating subjects in the control group and \(A=1\) indicating subjects in the treatment group. Under this setting, \(({X}_{5},\ldots ,{X}_{10})\) are considered as noisy biomarkers, while \({X}_{1}\) and \({X}_{2}\) are regarded as predictive biomarkers. On the other hand, biomarkers \({X}_{3}\) and \({X}_{4}\) are considered as prognostic biomarkers. The constants β₀ and β are the strength of prognostic and predictive effects, respectively.

For binary outcomes, we simulate the response \({Y}_{i}^{b}(A)\) ∼ Bernoulli (plogis(Y_i (A))), and \({Y}_{i}^{b}(A)\) can be expressed as:

where plogis\((x)=\frac{1\;+\;\tanh (x/2)}{2}\), and all other settings remained consistent with those used for continuous outcomes.

Simulation scenario II

In this scenario, we consider a simulation design with quadratic functions. The outcome

\(Y\) is simulated from a nonlinear model:

We simulate covariates X and error terms as in Simulation Scenario I, and the binary outcomes are simulated using the same method as in Simulation Scenario I.

Simulation scenario III

In this design, we incorporate interactive terms within the indicator function, adding more complexity to the simulation. To simulate data for a continuous outcome, we use the following model:

All other settings remain consistent with those used in Simulation Scenarios I and II.

Baseline models

We consider four baseline models in our analysis. The causal forest (CF) is a nonparametric model that extends the classic random forest algorithm to estimate conditional average treatment Effects. It is implemented via R package “grf”²⁹. The XGBoost with modified loss function (XGboostML) model is an adaptation of the XGBoost algorithm that integrates the A-learning loss function. It is implemented via our published R package “BioPred”³⁰. We also employ two linear regression models: linear regression with modified outcomes³¹ (LRMO) and linear regression with modified covariates³² (LRMC), both utilizing Lasso regularization³³. Of note, LRMO is only designed for continuous outcomes, while LRMC is suitable for binary outcomes. Both linear regression models are implemented by R package “glmnet”³⁴ with lasso penalty.

Evaluating the performance of models

In the simulation settings where the ground truth is known, we consider the true treatment benefiting group for patients with \(\{{i|}{Y}_{i}\left(1\right) > {Y}_{i}\left(0\right)\}\). Thus, an individual’s label is set to 1 if Y_i(1)>Y_i (0) and to 0 otherwise. When we assess the model’s performance in subgroup identification, we use evaluation metric of area under the ROC Curve (AUC) for classifying true subgroup labels. Concretely, our model produces a score \(\hat{f}\left(X\right)\) that reflects the likelihood of a patient belonging to the treatment-benefiting subgroup. By varying the threshold on \(\hat{f}\left(X\right)\), we obtain distinct pairs of true positive and false positive rates, forming an ROC curve. We then compute the AUC by integrating under this curve, in line with the evaluation strategies outlined in prior studies^7,19. Since XGboostML also incorporates an A-learning function into its loss, it follows the same approach as DeepRAB for computing the AUC. For other subgroup identification methods, the output is an estimated treatment effect. To compute the AUC in a similar way, we vary the decision boundary (i.e., the cutoff on the estimated treatment effect) and record the corresponding true positive and false positive rates against the known subgroup labels, thereby producing an ROC curve and an associated AUC.

To evaluate the model’s efficacy in identifying individual biomarkers, we rank the importance of each selected biomarker for each method. Specifically, for the tree-based models, CF and XGboostML, we derive importance scores using the variable_importance() function in the grf package and the predictive_biomarker_imp() function in the BioPred package. In the case of LRMO and LRMC, we base each biomarker’s ranking on the absolute magnitude of its corresponding coefficient. We consider the top two ranking biomarkers as identified biomarkers. We define detection rate as the frequency of each biomarker being chosen as one of the top two important features across \(N\) replications. Moreover, when we examine the model’s ability in detecting interactive biomarkers in simulation settings, we consider the model successfully picks out the interactive biomarker when two true biomarkers are detected as top two ranked biomarkers simultaneously for each replication.

Statistics and reproducibility

To optimize the performance of the DeepRAB, we have experimented with various tuning parameters including the learning rate (\(\eta\)), the number of layers (\(h\)) in decode layer, the dropout rate \(\delta\), the activation function ϕ_j, and the number of nodes (\({n}_{i},i=1,\ldots ,h\)) in each layer.

Optimizing these parameters solely based on the training dataset often leads to overfitting, where the algorithm performs well on the data but fails to generalize to other datasets. This issue is particularly pronounced when working with clinical datasets, which typically have small sample sizes. Therefore, it is a common practice in many machine learning algorithms to select values of tuning parameters using an independent validation dataset. When a validation dataset is not available, as is often the case in smaller clinical trials, cross-validation methods are employed. To address the overfitting problem, we have divided the entire dataset into training and test sets in an 8:2 ratio. We advocate determining the tuning parameters via \(K\)-fold cross-validation (CV) on the training data, recommending \(K=10\) for practical applications. The test set is reserved for final model evaluation after the optimal tuning parameters have been selected. The procedure for deriving the optimal tuning parameters is as follows: First, the training dataset is randomly split into 10 folds. DeepRAB is then trained on \(K\)-1 folds using different combinations of tuning parameters, and the error (A-learning loss) is estimated on the left-out fold. We calculate the average estimated errors on the left-out folds for each combination of tuning parameters. The best tuning parameters are identified through a grid search based on the smallest average errors.

Once the optimal tuning parameters are determined, we first re-train the entire training set (using both \(X\) and \(Y\)) to identify the predictive biomarkers. Next, we input the biomarkers \(X\) of the testing dataset into the optimized DeepRAB model to predict subgroup labels. In the simulation study, where the ground truth is known, we evaluate model performance using the AUC for classifying true subgroup labels. This tuning process ensures that DeepRAB is well-calibrated and minimizes the risk of overfitting, thereby enhancing its applicability in clinical settings. The detailed cross-validation procedure is illustrated in Fig. 3. Specifically, the Adam optimizer is utilized, \(\eta\) is selected from the set {0.01, 0.05, 0.001, 0.005}, \(h\) is chosen from {1, 2, 3, 4, 5}, \({n}_{i}\) is selected from {4, 8, 16, 32, 64, 128}, \(\delta\) is drawn from {0.2, 0.4, 0.6}, and \({\phi }_{j}\) is chosen from ReLU, Leaky ReLU, Sigmoid and Tanh. In addition, we set the pre-fixed number \(k\) equal to the number of covariates. Regarding the initialization of the temperature (\(T\)), we followed the approach proposed by Abid et al.¹⁸ to ensure effective exploration of different feature combinations and avoid convergence to suboptimal solutions. We initialize \(T\) as \(T(e)=\) \({T}_{1}({T}_{2}/{T}_{1})\)^e/E, where T (e) represents the temperature at epoch number \(e\), and \(E\) denotes the total number of epochs used for training the model. T₁ and T₂ are tuning parameters, typically set to high and low values, respectively. All other baseline models, along with their associated hyper-parameters, are optimized using the same approach. Each simulation scenario is repeated 1000 times to ensure a robust and reliable evaluation of model performance.

For the adalimumab dataset, the optimal tuning parameters for DeepRAB were determined as follows: the learning rate (\(\eta\)) was set to 0.001, with two hidden layers (\(h=2\)). The first hidden layer contained \({n}_{1}=\)16 nodes, and the second hidden layer contained \({n}_{2}=\)8 nodes. Leaky ReLU was used as the activation function for each hidden layer, and a dropout rate (\(\delta\)) of 0.2 was applied to prevent overfitting. Training was initiated with 10 to 100 epochs, and validation loss was continuously monitored. Early stopping was also implemented to prevent overfitting.

Reporting summary

Further information on research design is available in the Nature Portfolio Reporting Summary linked to this article.