Deep Semi-Supervised Survival Analysis for Predicting Cancer Prognosis

TL;DR

This work tackles limited labeled time-to-event data in cancer prognosis by applying deep semi-supervised learning to survival analysis. It introduces Cox-MT, a Mean Teacher-based neural Cox model that learns from both labeled and unlabeled data, combining a supervised negative partial likelihood with a consistency loss across teacher and student networks. The authors demonstrate strong gains for single-modal RNA-seq and WSI inputs, and further improvements with a multi-modal fusion using mutual attention to combine gene expression and imaging features. The approach reduces the need for large labeled datasets and shows robust improvements on TCGA cancers, suggesting broad applicability to other domains with censoring and scarce labels.

Abstract

The Cox Proportional Hazards (PH) model is widely used in survival analysis. Recently, artificial neural network (ANN)-based Cox-PH models have been developed. However, training these Cox models with high-dimensional features typically requires a substantial number of labeled samples containing information about time-to-event. The limited availability of labeled data for training often constrains the performance of ANN-based Cox models. To address this issue, we employed a deep semi-supervised learning (DSSL) approach to develop single- and multi-modal ANN-based Cox models based on the Mean Teacher (MT) framework, which utilizes both labeled and unlabeled data for training. We applied our model, named Cox-MT, to predict the prognosis of several types of cancer using data from The Cancer Genome Atlas (TCGA). Our single-modal Cox-MT models, utilizing TCGA RNA-seq data or whole slide images, significantly outperformed the existing ANN-based Cox model, Cox-nnet, using the same data set across four types of cancer considered. As the number of unlabeled samples increased, the performance of Cox-MT significantly improved with a given set of labeled data. Furthermore, our multi-modal Cox-MT model demonstrated considerably better performance than the single-modal model. In summary, the Cox-MT model effectively leverages both labeled and unlabeled data to significantly enhance prediction accuracy compared to existing ANN-based Cox models trained solely on labeled data.

Figures (5)

• Figure 1: Performance of single-modal Cox-MT and Cox-nnet models using gene expression values from TCGA to predict the prognosis of four types of cancer. a. Number of samples for each type of cancer. b. Average c-indexes of two models. c. Average IBSs. d. Box-plots of c-indexes. The Wilcoxon rank-sum test was used to obtain p-values. e. Box-pots of IBSs. f. KM curves for the two groups of patients in the test data stratified by the median HR of the training data predicted by the Cox-MT model: a high risk group (red) characterized by HRs exceeding the median and a low risk group (blue) defined by HRs below the median.
• Figure 2: Performance comparison for single-modal Cox-MT models using gene expression values or WSIs and the multi-modal Cox-MT model using both gene expression values and WSIs of TCGA breast cancer patients. The x-label MM stands for the multi-modal model. The Wilcoxon rank-sum test was used to obtain p-values.
• Figure 3: Performance of the single-modal Cox-MT model using gene expression values of breast cancer patients with various number of unlabeled data samples. TCGA+$n$, $n=1,000, 2,000, 3,000, 3,409$, stands for a data set that includes TCGA BRCA RNA-seq data and additional $n$ unlabeled breast tumors from the GSE96058 data set.
• Figure 4: Validation errors (1 - c-index) on the TCGA BRCA and GEO GSE96058 RNA-seq data over twenty runs per hyperparameter setting and their means. In each experiment, we varied one hyperparameter, and used default values for the rest.
• Figure 5: Cox-MT models for predicting cancer prognosis.a. Single-model Cox-MT model. b. The student model of the multi-modal Cox-MT model which has the same overall structure as the single-modal Cox-MT model.