Uncertainty Estimation in Deep Bayesian Survival Models

Bayesian methods can express uncertainty about their predictions, but have 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 neural networks. In this paper, we introduce the use of Bayesian inference techniques for survival analysis in neural networks that rely on the Cox 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. Our comprehensive experiments show that on the WHAS500 dataset, Bayesian techniques improve predictive performance over the state-of-the-art neural networks and on the larger SEER and SUPPORT datasets provide comparable performance. In all experiments, training with Monte Carlo Dropout is significantly faster than training with variational inference. 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.