SurvRNC: Learning Ordered Representations for Survival Prediction using Rank-N-Contrast

TL;DR

This study proposes Survival Rank-N Contrast (SurvRNC) method, which introduces a loss function as a regularizer to obtain an ordered representation based on the survival times, and demonstrates that using the SurvRNC method for training can achieve higher performance on different deep survival models.

Abstract

Predicting the likelihood of survival is of paramount importance for individuals diagnosed with cancer as it provides invaluable information regarding prognosis at an early stage. This knowledge enables the formulation of effective treatment plans that lead to improved patient outcomes. In the past few years, deep learning models have provided a feasible solution for assessing medical images, electronic health records, and genomic data to estimate cancer risk scores. However, these models often fall short of their potential because they struggle to learn regression-aware feature representations. In this study, we propose Survival Rank-N Contrast (SurvRNC) method, which introduces a loss function as a regularizer to obtain an ordered representation based on the survival times. This function can handle censored data and can be incorporated into any survival model to ensure that the learned representation is ordinal. The model was extensively evaluated on a HEad \& NeCK TumOR (HECKTOR) segmentation and the outcome-prediction task dataset. We demonstrate that using the SurvRNC method for training can achieve higher performance on different deep survival models. Additionally, it outperforms state-of-the-art methods by 3.6% on the concordance index. The code is publicly available on https://github.com/numanai/SurvRNC

Figures (3)

• Figure 1: An illustration of the deep neural network architecture trained using its native loss $\mathcal{L}_{Prognosis}$ and $\mathcal{L}_{SurvRNC}$.
• Figure 2: Overview of the proposed $\mathcal{L}_{SurvRNC}$ loss function for learning ordinal representations. In a randomly weighted sampled batch of $M$ patients, the loss function ranks them with respect to their time-to-event differences with the anchor. Contrasting the anchor patient with a positive pair patient enforces the similarity in the embedding space to be higher than the negative pair(s) of patients with a larger time-to-event difference than the positive pair. The uncertain patient pair(s), whose real-time difference with the anchor patient is unknown, are given less weight. (a) shows an example of an uncensored anchor and positive-pair patient. (b) shows an example of a censored anchor with an uncensored, positive pair. (c) provides all different combinations that can occur between a batch of patients.
• Figure 3: UMAP Visualization:$\mathcal{L}_{SurvRNC}$ with the native loss function of a DeepMTLR survival model on the SUPPORT dataset knaus1995support shows a better continuous latent feature representation (a) as compared to using only the native loss function (b).