Contrastive Learning for Predicting Cancer Prognosis Using Gene Expression Values

Several artificial neural networks (ANNs) have been developed recently to predict the prognosis of different types of cancer based on the tumor transcriptome. However, they have not demonstrated significantly better performance than the regularized Cox proportional hazards regression model. Training an ANN is challenging with a limited number of data samples and a high-dimensional feature space. Recent advancements in image classification have shown that contrastive learning (CL) can facilitate further learning tasks by learning good feature representation from a limited number of data samples. In this paper, we applied supervised CL to tumor gene expression and clinical data to learn feature representations in a low-dimensional space. We then used these learned features to train a Cox model for predicting cancer prognosis. Using data from The Cancer Genome Atlas (TCGA), we demonstrated that our CL-based Cox model (CLCox) significantly outperformed existing methods in predicting the prognosis of 19 types of cancer considered. We also developed CL-based classifiers to classify tumors into different risk groups and showed that CL can significantly improve classification accuracy. Specifically, our CL-based classifiers achieved an area under the receiver operating characteristic curve (AUC) of greater than 0.8 for 14 types of cancer and and an AUC of greater than 0.9 for 2 types of cancer. CLCox models and CL-based classifiers trained with TCGA lung cancer and prostate cancer data were validated with the data of two independent cohorts.