Uncertainty Estimation in Deep Bayesian Survival Models

Christian Marius Lillelund; Martin Magris; Christian Fischer Pedersen

Abstract

Bayesian methods can express uncertainty about their predictions, but has seen little adaptation in survival analysis using neural networks. Proper uncertainty estimation is important in high-risk domains, such as the healthcare or medical field, if machine learning methods are to be adopted for decision-making purposes, however uncertainty estimation is a known shortcoming of NNs. In this paper, we introduce the use of Bayesian inference techniques for survival analysis in neural networks that rely on the Cox's proportional hazard assumption, for which we discuss a new flexible and effective architecture. We implement three architectures: a fully-deterministic neural network that acts as a baseline, a Bayesian model using variational inference and one using Monte-Carlo Dropout. We show with comprehensive experiments that the Bayesian models improve predictive performance over SOTA neural networks in a test dataset with few samples (WHAS500, 500 samples) and provide comparable performance in two larger ones (SEER and SUPPORT, 4024 and 8873 samples, respectively), however using variational inference comes with longer training times. Our Bayesian models additionally provide quantification of both aleatoric and epistemic uncertainty, which we exhibit by plotting 95\% confidence intervals around the survival function and showing a probability density function of the survival time. Our work motivates further work in leveraging uncertainty for survival analysis using neural networks.

Keywords: Uncertainty estimation; Neural networks; Survival analysis; Variational inference; MC Dropout

Links
[Full text PDF][Bibtex]