ContraSurv: Enhancing Prognostic Assessment of Medical Images via Data-Efficient Weakly Supervised Contrastive Learning.

Abstract

Prognostic assessment remains a critical challenge in medical research, often limited by the lack of well-labeled data. In this work, we introduce ContraSurv, a weakly-supervised learning framework based on contrastive learning, designed to enhance prognostic predictions in 3D medical images. ContraSurv utilizes both the self-supervised information inherent in unlabeled data and the weakly-supervised cues present in censored data, refining its capacity to extract prognostic representations. For this purpose, we establish a Vision Transformer architecture optimized for our medical image datasets and introduce novel methodologies for both self-supervised and supervised contrastive learning for prognostic assessment. Additionally, we propose a specialized supervised contrastive loss function and introduce SurvMix, a novel data augmentation technique for survival analysis. Evaluations were conducted across three cancer types and two imaging modalities on three real-world datasets. The results confirmed the enhanced performance of ContraSurv over competing methods, particularly in data with a high censoring rate.