Uncertainty Estimation in Deep Bayesian Survival Models

TY - GEN

T1 - Uncertainty Estimation in Deep Bayesian Survival Models

AU - Lillelund, Christian Marius

AU - Magris, Martin

AU - Pedersen, Christian Fischer

PY - 2023/10

Y1 - 2023/10

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

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

KW - MC Dropout

KW - neural networks

KW - survival analysis

KW - uncertainty estimation

KW - variational inference

UR - http://www.scopus.com/inward/record.url?scp=85179522978&partnerID=8YFLogxK

U2 - 10.1109/BHI58575.2023.10313466

DO - 10.1109/BHI58575.2023.10313466

M3 - Article in proceedings

T3 - IEEE EMBS International Conference on Biomedical and Health Informatics (BHI)

BT - IEEE-EMBS International Conference on Biomedical and Health Informatics (BHI’23)

PB - IEEE

T2 - IEEE-EMBS International Conference on Biomedical and Health Informatics

Y2 - 15 October 2023 through 18 October 2023

ER -