DWCox: A density-weighted Cox model for outlier-robust prediction of prostate cancer survival

Feature construction

Training DWCox requires a training group of N patients whose clinical features X and survival outcomes Y are known. X is an N-by-M matrix, where M is the number of clinical features and each element X_ij is the value of the jth clinical feature of the ith patient. Y is an N-by-2 matrix, where each row gives the survival outcome of a patient. The 1st column of Y is a vector of the last observed survival time t, and the 2nd column is a vector of binary event indicators d. A patient i with d_i= TRUE is known to die at time t_i. Oppositely, one with d_i= FALSE is known to be alive at time t_i, but no information is available after t_i. In the latter case, the record of that patient is said to be censored. In the data sets used in the PCDC, Y is known, while X needs to be constructed from clinical data.

To ensure fair comparison with Halabi’s model, DWCox constructed X in line with the way Halabi defined his 22 features, as summarized in Table 1 and described in details in the Supplementary material. Note that two features Halabi’s model started with, namely the Charlson comorbidity index and the Biopsy Gleason score, were not considered by DWCox since during the PCDC the former was not available in the training data and the latter was 100% missing in the leaderboard data. (Data were split into training, leaderboard and final validation sets. Details will be described in the Experiments section). That means M = 20.

At this stage X was not complete (i.e. there were many missing elements in that matrix) due to missing information in the raw clinical records. Those missing values in X were imputed with the algorithm Multivariate Imputation by Chained Equations (MICE)^5,6. The idea of MICE is to use Bayesian statistics to iteratively infer the missing values from other known and previously inferred values. Missing values in the training data were imputed with knowledge about the survival outcome, since it was argued that the outcomes could help generate less biased imputations⁷. The survival outcome was incorporated into the imputation in the form of the Nelson–Aalen estimator as suggested by White and Royston⁸. Imputation on the leaderboard and final validation data were done without using the survival outcome.

During the PCDC, three more binary features were used to indicate the trial ID (described in the Experiments section) of each patient. Those features were removed in post-challenge analysis so that the performance of DWCox does not depend on prior knowledge about the data source.

Density-based weighting

After the imputation, the N-by-M matrix X can be represented by N points scattered in a M-dimensional space 𝔽 (“feature space”). Each point represents a patient whose each coordinate is the value of one of his/her M clinical features. We assign each patient i a weight w_i ∈ [0, 1] proportional to the estimated local Gaussian kernel density in the feature space. To calculate w_i, we used the default settings of the function kepdf in the R package pdfCluster⁹. These weights were then divided by the maximum value. Thus a patient with a higher weight indicates there are more other patients with similar clinical features.

Model training

After density-based weighting, we used the R package glmnet¹⁰ to maximize the weighted partial likelihood

During the PCDC, L₂ regularization was imposed to the objective function. The penalty weight was chosen to optimize the model performance (more specifically, iAUC, as defined in the next subsection) averaged over 100 repeated random sub-sampling validation on the training data. In each random sub-sampling validation experiment, 2/3 of all the training data were randomly selected to train the model with a wide range of possible penalty weights, and the iAUC was evaluated for each possible penalty weight on the remaining 1/3 of the training data. After the PCDC, the regularization was removed from DWCox since its contribution to the model performance was not obvious during the challenge and its removal sped up training.

After model training, the risk vector r of the training patients were calculated as

Prediction & evaluation

The trained model was used to predict the risk r_test and the remaining lifetime t_test for a new group of patients whose clinical features X_test could be constructed from clinical data while the outcome Y_test was not seen by the model. The model performance was then evaluated by comparing r_test and t_test to Y_test.

The predicted risks r_test were evaluated with the integrated area under the ROC curve (iAUC) as described below. After obtaining $\widehat{\mathit{\beta }}$ by maximizing Equation (1), we can estimate the risks of the patients ${\mathbf{\text{r}}}_{\text{test}}={\mathbf{\text{X}}}_{\text{test}}\widehat{\mathit{\beta }}.$ Then an estimated order of death ô can be constructed by sorting r_test (i.e. ô_i = j where i = 1, 2, … , N and r_test,i is the jth smallest element of r_test). By comparing ô with the actual outcome Y_test, at any given time threshold t_i we can calculate the area under the receiver operating characteristic curve AUC_ti. If we integrate AUC_ti with respect to t_i from the 6th to the 30th month, we get the integrated area under curve iAUC ∈ [0, 1]. The greater the iAUC, the better the predicted risks reflect the actual order of death.

DWCox also gives the estimated time to death of the test set: ${\widehat{\mathbf{\text{t}}}}_{\text{test}}=\widehat{k}{\mathbf{\text{r}}}_{\text{test}}+\widehat{\mathbf{\text{b}}}.$ In the PCDC ${\widehat{\text{t}}}_{\text{test}}$ was evaluated by its RMSE from t_test.

Extended applications

The open-source release of DWCox is coded in a way such that it can be easily re-trained and applied to other survival analysis problems, not restricted to prostate cancer. To re-train and apply DWCox to a new dataset, users simply need to:

Here X and X_test can have as many rows (i.e. subjects) and columns (i.e. features) as needed. They can have missing values as well. More details can be found in the documentation inside the package.

Challenge data & context

DWCox has been developed and evaluated with data from the comparator arms of four phase III clinical trials with over 2,000 mCRPC patients in total treated with first-line docetaxel. Those four trials and the corresponding data providers are:

During the PCDC those trials were referred to with their study IDs (Table 2).

The development and evaluation of DWCox began with the 2015 Prostate Cancer DREAM Challenge and continued after the challenge. The full anonymized information about the patients in trials ASCENT-2, MAINSAIL and VENICE was released to the challenge participants. As for trial ENTHUSE-33, the participants only knew the clinical records available at the beginning of the trial ("baseline clinical records"), while data obtained after the start of the trial including the survival outcome were visible only to the challenge organizers. The challenge goal was to develop models that used the baseline clinical records to predict the patients’ relative risk (sub-challenge 1a), days till death (sub-challenge 1b), and treatment discontinuation (sub-challenge 2) (Table 3).

DWCox was trained on Trials ASCENT-2, MAINSAIL and VENICE (“PCDC training data”) by the authors, and evaluated on Trial ENTHUSE-33 (“PCDC validation data”) by the challenge organizers. Trial ENTHUSE-33 was further divided into a leaderboard set (157 patients) and a validation set (313 patients). The leaderboard set was used to run three leaderboard rounds. In each round, the challenge organizers randomly subsampled 80% patients from the leaderboard set, evaluated the participants’ models on that random sample, and returned the feedback to the participants. After the 3rd leaderboard round, each participating team submitted a final model, whose performance on the validation set was used to rank the teams. Bootstrapping was performed by the challenge organizers to make sure the winning teams gave statistically significantly better predictions than other teams and Halabi’s model. DWCox was involved in the leaderboard rounds of sub-challenge 1a and the final scoring round of sub-challenges 1a & 1b.

Simulations

Simulation experiments were performed to evaluate the contribution of density-based weighting to the model performance. DWCox was trained and evaluated on 100 simulated data sets (one example is given in Figure 2) separately, each of which was designed to mimic the real challenge data to some extent, while the randomness in the data generation process assured the variation across simulations. In each simulation, three groups of patients were simulated. Each patient had 20 features and an outcome.

Figure 2. Scatter plot of the first two principle components of the signal, noise and validation groups in a simulated data set.

Each point represents a patient. The shapes mark the mean of each group. (Best viewed in color).

One group (“signal group”) represented a group of 1,000 patients that reflected the true correlation between the outcome and the features. The features were sampled from Gaussian distributions:

We would like to clarify a few things about Equation (4). Readers may get confused if they see an online manuscript with the same title and authors as those of Reference 15, where the minus sign of Equation (4) is outside the parenthesis. Obviously it is a typo, and it has been corrected in the version cited here. Although Equation (4) may not look like a Weibull distribution at first glance, the proof is a very straightforward and standard procedure. The shape and scale parameters of the Weibull distribution is ν and ${\left(\lambda e^{{\mathbf{\text{X}}}_{\text{signal},i*}^T\mathit{\beta }}\right)}^{-1/\nu }$ respectively.

Such generated survival times follow a Cox model with the baseline hazard function h₀(t) = λνt^ν−115. The parameters λ, ν and β were estimated from the uncensored part of the PCDC training data as follows. First, we assumed β = 0 and fit a Weibull distribution to the distribution of t_uncensored to estimate ν and λ. Then DWCox was applied to the PCDC training data to obtain $\widehat{\mathit{\beta }}$. At this stage $\widehat{\mathit{\beta }}$ did not include ${\widehat{\beta }}_0$, a constant term that affected $\widehat{\mathbf{\text{t}}}$ but not iAUC, since ${\widehat{\beta }}_0$ played no role during the maximization process of Equation (1). We chose a ${\widehat{\beta }}_0$ value such that the mean of the survival times simulated with Equation (4) was close to the mean of the uncensored survival times in the PCDC training data. After getting the estimates of λ, ν and β, t_signal was simulated with Equation (4).

We then generated 1,000 more patients (“noise group”) to represent outliers, or noises, in the training data. We made the outliers more sparsely and widely distributed in the feature space than the signal group by simulating

The survival times of the noise group were simulated with a Weibull distribution independent of X_noise:

A 3rd group of 500 patients (“validation group”) was generated in a fashion similar to that of the signal group.

We let

After simulating the three groups of patients, we mixed the signal and noise groups together to form a training set. DWCox and a 20-feature standard Cox model were trained on this training set, and evaluated with iAUC on the validation group.

Heterogeneity in the PCDC data

Analysis of the PCDC data suggests that there exists rather high heterogeneity across the three training trials and the validation trial. The missing-rate profile of the 20 clinical features varies across trials (Figure 3). The average values of the first two principle components of the 20 features of Trial ASCENT-2 is farther away from those of the validation trial, compared to those of the other two training trials (Figure 4). Leave-one-trial-out cross-validation (i.e. to train with two training trials and evaluate with the left-out training trial) gives very different results when different trials are left out (Table 4).

Figure 3. Heatmap of the percentage missing of the 20 clinical features used in DWCox.

Figure 4. Scatter plot of the first two principle components of the four prostate cancer trials.

Each point represents a patient. The shapes mark the average values of each trial.

Those facts give such a clue: If we consider the “true model” underlying the validation trial as the signal, it is very likely that the PCDC training data contain many outliers. Those outliers do not follow the “true model”, and thus tend to bring down the validation-set performance of models that failed to deal with the outliers properly during training. Therefore robustness against outliers is probably important to models aimed at winning the PCDC.

Indeed, several other winning teams of the PCDC tried hard to deal with the outliers in the training data. For example, the top performer (FIMM-UTU) of sub-challenge 1a decided to discard the entire ASCENT-2 trial, because after some manual exploration in the data space they found significant differences in clinical variables that set this trial apart from the other trials. Our team (Team Cornfield) instead used all available data and let DWCox automatically handle the outliers.

Results on the PCDC data

DWCox obtained the best average ranking in sub-challenges 1a & 1b among about 50 models (Figure 5). On the PCDC validation data, DWCox gave an iAUC of 0.7789 and a RMSE of 194.8650 days, out-performing Halabi’s model which gave an iAUC of 0.7581 and a RMSE of 196.6704 days. Bootstrapping has shown that DWCox outperforms Halabi’s model with a Bayes Factor (BF) > 3. Note that while the other numbers in this paragraph are official results provided by the challenge organizers, the Halabi RMSE is not. In order to get the Halabi RMSE, we implemented a Halabi’s model and appended to it a linear regression step similar to the one in DWCox. After applying bootstrapping and the BF > 3 threshold against other teams’ submissions, the challenge organizers reported DWCox as a winner in sub-challenge 1b and a runner-up in sub-challenge 1a. The winner of sub-challenge 1a, FIMM-UTU, obtained an iAUC of 0.7915 and a RMSE of 201.3779 days. Their model is an ensemble of penalized cox regressions developed with extensive manual exploration in the data space. More details about the challenge results can be found at https://www.synapse.org/#!Synapse:syn2813558/wiki/232674. Table 1 gives the regression coefficients determined by DWCox.

Figure 5. Ranking of the top teams in sub-challenges 1a & 1b.

The six best teams of each sub-challenge are included. DWCox was submitted by the authors’ Team Cornfield.

An inverse correlation between the actual survival time t and risk scores r was observed (Figure 6). Note that the adjusted R² of the linear regression $\widehat{\mathbf{t}}=\widehat{k}\mathbf{r}+\widehat{\mathbf{b}}$ is small (0.1513), and the shape of the t vs r plot implies that there may exist models better than a linear regression for capturing their correlation.

Figure 6. Scatter plot of the uncensored survival time versus the predicted risk on the PCDC training data.

The straight line is the linear regression line with slope = -234.6, intercept = 810.3 and adjusted R² = 0.1513.

Results on simulated data

In the 100 repeated simulations (described in the Experiments section), DWCox performed better than a standard Cox model when as many as half of the training data were outliers. DWCox not only gave better average performance over the 100 experiments (Table 5, Figure 7), but also performed consistently better in each experiment (Figure 8, paired t-test p-value = 2.1 × 10⁻²⁰). The improvement in performance clearly resulted from the density-based weighting, since everything else was the same across the two models.

Figure 7. Boxplot of the iAUC of DWCox and a standard Cox model in 100 simulations.

The boxes show the medians and inter-quartile ranges (IQR). The vertical black lines extends from the boxes by at most 1.5 IQR. Black points represent experiments whose iAUC is more than 1.5 IQR away from the boxes.

Figure 8. DWCox iAUC vs the standard Cox iAUC in 100 simulations.

Each point is given by a simulation. The straight line has slope = 1 and intercept = 0.

Note that in the simulations we used iAUC but not the RMSE to evaluate model performance. There are three reasons for that. 1. iAUC evaluates model performance on the validation data in a more comprehensive manner, while RMSE is based on individual predictions which are independent of each other. 2. DWCox’s time-to-event prediction is dependent on its predicted risks. 3. A standard Cox model does not directly give the predicted time-to-event.