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

Dongjoon Lee, Hyeryn Park, Changhee Lee

Keywords: Survival Analysis, Contrastive Learning, Deep Learning, Healthcare

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 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.

Primary Area: Machine learning for healthcare

Submission Number: 15583