Interpretable deep learning for improving cancer patient survival based on personal transcriptomes

TCGA cancer genomics and clinical data

To predict the survival time of cancer patients based on their transcriptome profiles and drug treatments, we obtained the transcriptome and clinical data from TCGA. A total of 4311 patients have both transcriptome data and clinical information such as drug prescriptions and survival time, covering 178 drugs and 33 cancer types. The most abundant samples are from breast and ovarian tumors (Fig. 1a), and the gender distribution is 63.4% females and 36.6% males (Fig. 1b). Based on the vital status (Fig. 1c), 1645 patients (38.2%) were deceased. The distribution of their days-to-death is shown in Fig. 1d. For the other 2666 (61.8%) alive or censored patients, we used the days-to-the-last-follow-up (distribution shown in Fig. 1e) to infer whether a patient was long-lived. We trained a Gene Ontology (GO)-guided deep learning network to predict the impact of different drugs on the survival time given the patient’s gene expression profile. The actual number of patients used in our network training was slightly less, depending on whether the Morgan fingerprint of the drug is available.

Two-stage training with knowledge transfer to fully utilize data

Transfer learning was initially devised to transfer the learner trained in similar but different domains to the target domain²⁰. Recent studies show that two-stage training can improve deep learning models in the data-scarce setting^21,22. Here the majority of available TCGA data are censored. Therefore, to harness the power of machine learning in personalized medicine, we devised a two-stage training schema in our CancerIDP with knowledge transfer.

In the first stage, the network weights are randomly initialized, and we aim to utilize as many patients as training samples as possible. To do that, we included both deceased and qualified censored patients to create binary classification supervision signals: long-lived or short-lived. Positive samples are long-lived patients known to be alive after 1,200 days, including patients with either days-to-death or days-to-the-last-follow-up greater than 1200. Negative samples are short-lived patients who died within 1000 days (i.e., days-to-death < 1000). The first stage’s goal is to exploit as much data as possible to abstract deep representations from gene expression and drug structures for the vital status prediction (more straightforward binary classification), which will be used as a starting point in the second stage for fine-grind survival time prediction.

In the second stage, the network structure remains the same (Fig. 2), and the weights are initialized as those in the best-performing model on the validation set from the first stage. The assumption is that the feature extractor learned in the binary survival status classification can be a good starting point for a more fine-grind survival time prediction. Only deceased patients with exact days-to-death records are used in the second stage. The binary network head from the first stage is discarded, and a new regression head is trained from scratch to predict the days-to-death given a patient’s gene expression profile and the prescribed drug. In the actual training, the days-to-death was transformed to months-to-death on a log scale. Figure 3a shows that the best-performing model achieves 96% classification accuracy in the first stage on the hold-out data. In the second stage, Pearson’s R equals 0.937 between the ground truth months-to-death and the prediction (Fig. 3b).

The memorization pitfall when predicting survival responses from drug treatments

Deep learning models have impressive expressiveness, but in practice, we noticed a potential pitfall that comes with the ability of the deep network to emulate any function. We called this the “memorization effect” in our survival prediction. In our earlier experiments, we included all drugs in the deep learning model and prepared the training and the hold-out validation sets by the random split. However, we observed that for some rarely prescribed drugs with only one or two supporting data points (i.e., 1–2 patients) in the training set, the network tends to simply “memorize” such drugs instead of generalizing meaningful embedding. And then, in the testing phase, the network just outputs the survival time of the patients that it has memorized (observed) in the training phase. Such effect is demonstrated in Fig. 3c: some of the predicted months-to-death values (log scale) are the same since the model only observed one patient with this specific drug during the training, and the model output the survival time of that patient in the testing when fed into this particular drug.

To counter the memorization effect, we eliminated drugs with less than five supporting data points in the training set. Although the number of training samples decreases slightly, the network does not suffer from the memorization effect but learns meaningful embeddings and is more generalizable. Pearson’s R improves from 0.834 to 0.937 after eradicating the memorization effect in the months-to-death prediction (Fig. 3c,d).

Improving patients’ survival based on personalized medicine enabled by CancerIDP

Distinct from other deep-learning-based cancer survival time predictions, our model not only predicts survival time but also enables optimal drug selection for each patient based on the expression encoder and the drug encoder. From the existing clinical data, we know that patient X with transcriptome T is prescribed drug D, and the survival time is S. This triplet (T, D, S) forms a training data point for our network. With the predictive ability achieved through training, the network has learned to encode transcriptomes and drugs as predictive deep embeddings. We can then predict the drug response in silico by pairing each patient with each medication and identifying the optimal drug (D*) with the longest survival time for patient X.

In this TCGA cancer genomics dataset, based on our in-silico drug selection, we can find a drug different from the actual prescription for 27.4% of patients, with which the predicted survival time is longer than the observed survival time (Fig. 4a). Figure 4b compares the actual survival time for the prescribed drugs and the predicted survival time if paired with the deeply-learned optimal drugs. Deep-learning-enabled personalized medicine can significantly improve patient survival by increasing the median of months-to-death from 27.1 to 31.0. Noticeably, 23.4% of patients can have a relative improvement in survival time by at least 5% (Fig. 4c, colored in orange, the size of the mark is proportional to the relative improvement). The most significant improvement is 2.97-fold, where the months-to-death is 24.7 for the prescribed drug and 73.4 for the optimal drug suggested in-silico.

To support our in-silico drug discovery, we randomly selected a patient (TCGA-ZH-A8Y4, patient 1 in Supplementary Fig. S1) and identified nine more patients with the most similar transcriptome profiling according to the cosine similarity between transcriptomes (normalized to unit vectors). Three of the ten patients had their predicted optimal drugs matching their actual prescribed drugs. Five drugs selected by CancerIDP were actually prescribed to at least one of the ten patients. Particularly, the drug Halichondrin B (canonical SMILE: CC1CC2CCC3C(=C)CC(O3)CCC45CC6C(O4)C7C(O6)C(O5)C8C(O7)CCC(O8)CC(=O)OC9C(C3C(CC4C(O3)CC3(O4)CC4C(O3)C(CC3(O4)CC(C4C(O3)CC(O4)C(CC(CO)O)O)C)C)OC9CC(C1=C)O2)C, PubChem CID: 44581725) was prescribed to patients 3 and 9, and predicted to be the optimal drug for patients 4, 5, and 9. After being prescribed Halichondrin B, patient 3 survived for 24.69 months, and patient 9 survived for 34.12 months. In contrast, patients 4 and 5 only survived for 18.21 and 29.45 months under different drug treatments. As predicted by CancerIDP, prescribing Halichondrin B could be beneficial for patients 4 and 5, as it may lead to longer survival times of 22.21 and 33.35 months, respectively, which is similar to its effect on patient 9. Thus, information obtained from CancerIDP could be used as input for an algorithmic expert system to aid clinical decision-making.

To further support our in-silico drug discovery, we also checked the newly approved cancer drugs since 2010 (summarized in²³). These newly approved drugs usually increase cancer patients’ survival time. Our CancerIDP model included 46 patients involving these recently approved drugs. Of these, 35 patients had predicted drugs matching their actual prescribed newly approved drugs. Our analysis using CancerIDP showed that prescribing these newly approved drugs could increase the survival times of seven patients who were receiving other drugs. In contrast, CancerIDP predicted other drugs for only four patients who were prescribed these new drugs. These results may suggest that our CancerIDP predictions are more aligned with the recommendations of medical companies. Thus, deep representation learning provides great promise here. As the system is fed with more abundant data, such data-driven personalized drug prescriptions will become more valuable and trustworthy.

Learn transcriptome embeddings that are predictive of months-to-death

The final prediction of the network is based on two components: the transcriptome embedding learned from the GO-guided encoder, and the drug embedding from the autoencoder. The drug autoencoder is simply a black-box model that learns representations from drug chemical structures expressed in the Morgan fingerprints. We are currently unable to interpret the mechanisms. However, our transcriptome embedding is based on the hierarchy structure of GO terms and has greater interpretability of how genes act together to predict drug survival responses.

We inspected the transcriptome embedding learned from all patients using t-SNE and visualized the top two principal components (Fig. 5a–c). The transcriptome embedding is informative about months-to-death, as shown in Fig. 5a. Patients with similar transcriptome embeddings share similar months-to-death values. The expression of some marker genes such as ESR1 (implicated in hormone resistance and anti-estrogen therapies for breast cancer, Fig. 5b) or KIFC3 (encodes a member of the kinesin-14 family of microtubule motors, Fig. 5c) is consistent with the spatial pattern of the transcriptome embedding, indicating the contribution of these genes to the transcriptome embedding for drug survival responses.

Discover the most discriminating GO components for a cancer patient

Cellular responses of drugs depend on both transcriptomes and drug chemical structures²⁴. In CancerIDP, we encode the expression information and drug information via transcriptome GO encoder and drug autoencoder (details in “Methods”). Since the neural network structure of the transcriptome GO encoder is designed according to the GO system structure in human cells, it is the key to the interpretability of CancerIDP. Specifically, we build the network based on the GO system. Each GO term (node) is represented by a multi-layer perceptron (MLP). These network nodes take inputs from child GO terms and send outputs to parent GO terms. The state of each GO term is represented by the activation of MLP, and the connectivity weights between nodes are determined through end-to-end network training. We applied the attention mechanism^25,26 to learn the most important GO terms for survival time prediction (details in “Methods”). Attention weight assigned to each GO node indicates the relative importance of that GO term to the prediction task at hand. After training, when we forward pass a patient’s gene expression profile, we interpret the nodes with the highest attention weights as the most distinguished GO terms for survival determination for that patient.

Figure 5d shows an example of patient TCGA-57-1582 with ovarian cancer. The most discriminating GO terms are highlighted in red, including catabolic processes like carbohydrate catabolic process (GO:0016052), macromolecule catabolic process (GO:0009057), T-cell co-stimulation (GO:0031295), and some signaling pathways involved in phagocytosis. Therefore, such a GO-guided encoder informs the activated pathway at the individual level, paving the way for the mechanistic study of personalized treatment.

Similarly, we forward passed each patient’s gene expression to the trained GO encoder, and obtained attention weights of GO terms for each patient. We then summarized the most common GO terms with large average attention weights across patients. The discriminating GO terms with average attention weights across patients ≥ 0.15 are provided in Supplementary Table S1. Their positions in the pathway hierarchy are shown in Supplementary Fig. S2. Important GO terms include T cell mediated immune response to tumor cell (GO:0002424), negative regulation of G2/M transition of mitotic cell cycle (GO:0010972), extrinsic apoptotic signaling pathway via death domain receptors (GO:0008625), and so on.