Toward a Well-Calibrated Discrimination via Survival Outcome-Aware Contrastive Learning

TL;DR

This work proposes a novel contrastive learning approach specifically designed to enhance discrimination performance without sacrificing calibration, which outperforms state-of-the-art deep survival models in both discrimination and calibration.

Abstract

Previous deep learning approaches for survival analysis have primarily relied on ranking losses to improve discrimination performance, which often comes at the expense of calibration performance. To address such an issue, we propose a novel contrastive learning approach specifically designed to enhance discrimination \textit{without} sacrificing calibration. Our method employs weighted sampling within a contrastive learning framework, assigning lower penalties to samples with similar survival outcomes. This aligns well with the assumption that patients with similar event times share similar clinical statuses. Consequently, when augmented with the commonly used negative log-likelihood loss, our approach significantly improves discrimination performance without directly manipulating the model outputs, thereby achieving better calibration. Experiments on multiple real-world clinical datasets demonstrate that our method outperforms state-of-the-art deep survival models in both discrimination and calibration. Through comprehensive ablation studies, we further validate the effectiveness of our approach through quantitative and qualitative analyses.

Figures (9)

• Figure 1: An illustration of the network architecture for ConSurv.
• Figure 2: t-SNE visualization for latent representations learned with $\mathcal{L}_{\text{NLL}}$ only, $\mathcal{L}_{\text{SNCE}}$ only, and ConSurv for the METABRIC dataset, colored by event times (for uncensored samples).
• Figure 3: Calibration plots for ConSurv in comparison with benchmarks for the METABRIC dataset.
• Figure 4: Comparison of the survival curves across various patient subgroups for the METABRIC dataset.
• Figure 5: t-SNE visualization for latent representations learned with $\mathcal{L}_{\text{NLL}}$ only, $\mathcal{L}_{\text{SNCE}}$ only, and ConSurv for the FLCHAIN and SUPPORT datasets, colored by event times (for uncensored samples).