Diffsurv: Differentiable sorting for censored time-to-event data.

Andre Vauvelle, Benjamin Wild, Aylin Cakiroglu, Roland Eils, Spiros Denaxas

Keywords: Survival analysis, Censored data, Semi-supervised learning, Time-to-event-data, Algorithmic Supervision

TL;DR: Differentiable sorting for censored survival analysis.

Abstract: Survival analysis is a crucial semi-supervised task in machine learning with numerous real-world applications, particularly in healthcare. Currently, the most common approach to survival analysis is based on Cox's partial likelihood, which can be interpreted as a ranking model optimized on a lower bound of the concordance index. However, this interpretation of Cox's partial likelihood considers only pairwise comparisons. Recent work has developed differentiable sorting methods which relax this pairwise independence assumption, enabling the ranking of sets of samples. However, current differentiable sorting methods can not account for censoring, a key factor in many real-world datasets. To address this limitation, we propose a novel method called \emph{diffsurv}. We extend differentiable sorting methods to handle censored survival analysis tasks by predicting matrices of possible permutations that take into account the uncertainty introduced by censored samples. We contrast this approach with methods derived from partial likelihood and ranking losses. Our experiments show that diffsurv outperforms established baselines in various simulated and real-world risk prediction scenarios. Additionally, we demonstrate the benefits of the algorithmic supervision enabled by diffsurv by presenting a novel method for top-k risk prediction that outperforms current methods.

Anonymity: We agree to keep the submission (including any supplements and/or code) anonymous.

Formatting: We confirm that we read and complied with the author's instructions.