Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

NeuralSurv: Deep Survival Analysis with Bayesian Uncertainty Quantification

Mélodie Monod, Alessandro Micheli, Samir Bhatt

TL;DR

NeuralSurv addresses the need for uncertainty-aware, high-capacity survival models in continuous time. It combines deep neural networks with Bayesian inference by introducing a two-stage data-augmentation strategy (Pólya–Gamma and marked Poisson processes) and a local linearization of the Bayesian NN to achieve conjugacy and scalable VI. The approach yields well-calibrated survival functions with credible intervals, outperforming state-of-the-art deep survival models in calibration while maintaining competitive discriminative performance across synthetic and real datasets. This framework enables robust uncertainty quantification in data-scarce regimes and sets a foundation for scalable, Bayesian deep survival analysis in complex settings.

Abstract

We introduce NeuralSurv, the first deep survival model to incorporate Bayesian uncertainty quantification. Our non-parametric, architecture-agnostic framework captures time-varying covariate-risk relationships in continuous time via a novel two-stage data-augmentation scheme, for which we establish theoretical guarantees. For efficient posterior inference, we introduce a mean-field variational algorithm with coordinate-ascent updates that scale linearly in model size. By locally linearizing the Bayesian neural network, we obtain full conjugacy and derive all coordinate updates in closed form. In experiments, NeuralSurv delivers superior calibration compared to state-of-the-art deep survival models, while matching or exceeding their discriminative performance across both synthetic benchmarks and real-world datasets. Our results demonstrate the value of Bayesian principles in data-scarce regimes by enhancing model calibration and providing robust, well-calibrated uncertainty estimates for the survival function.

NeuralSurv: Deep Survival Analysis with Bayesian Uncertainty Quantification

TL;DR

NeuralSurv addresses the need for uncertainty-aware, high-capacity survival models in continuous time. It combines deep neural networks with Bayesian inference by introducing a two-stage data-augmentation strategy (Pólya–Gamma and marked Poisson processes) and a local linearization of the Bayesian NN to achieve conjugacy and scalable VI. The approach yields well-calibrated survival functions with credible intervals, outperforming state-of-the-art deep survival models in calibration while maintaining competitive discriminative performance across synthetic and real datasets. This framework enables robust uncertainty quantification in data-scarce regimes and sets a foundation for scalable, Bayesian deep survival analysis in complex settings.

Abstract

We introduce NeuralSurv, the first deep survival model to incorporate Bayesian uncertainty quantification. Our non-parametric, architecture-agnostic framework captures time-varying covariate-risk relationships in continuous time via a novel two-stage data-augmentation scheme, for which we establish theoretical guarantees. For efficient posterior inference, we introduce a mean-field variational algorithm with coordinate-ascent updates that scale linearly in model size. By locally linearizing the Bayesian neural network, we obtain full conjugacy and derive all coordinate updates in closed form. In experiments, NeuralSurv delivers superior calibration compared to state-of-the-art deep survival models, while matching or exceeding their discriminative performance across both synthetic benchmarks and real-world datasets. Our results demonstrate the value of Bayesian principles in data-scarce regimes by enhancing model calibration and providing robust, well-calibrated uncertainty estimates for the survival function.
Paper Structure (111 sections, 10 theorems, 157 equations, 1 figure, 12 tables, 2 algorithms)

This paper contains 111 sections, 10 theorems, 157 equations, 1 figure, 12 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. NeuralSurv
  3. Sigmoidal Hazard Function
  4. Likelihood Distribution
  5. Prior Distributions
  6. Prior Distribution on $\boldsymbol{\theta}$.
  7. Prior Distribution on $\phi$.
  8. Normalization Factor $Z$.
  9. Posterior Distribution
  10. Data Augmentation Strategy
  11. Pólya-Gamma Augmentation Scheme
  12. Poisson Process Augmentation Scheme
  13. Augmented Likelihood
  14. Variational Inference in the Augmented Space
  15. Variational Mean–Field Approximation
  16. ...and 96 more sections

Key Result

Theorem 3.1

Assume for each $i=1,\ldots,N$ that the function $g(\cdot,\mathbf{x}_{i};\cdot)\in C([0,y_{i}]\times\mathbb{R}^{m})$. Let $p(y_i, \delta_i \mid \mathbf{x}_i, \phi, g(\cdot;\boldsymbol{\theta}))$ be the likelihood density given in eq-likelihood. Additionally, let $p\left( y_i, \delta_i\mid \mathbf{x}

Figures (1)

  • Figure 1: Comparison of the true survival function (black) with the estimated survival functions from NeuralSurv and the two top-performing benchmark models (colored) on synthetic data. The time axis is truncated at the maximum observed event time in the training data. Each panel represents a different training set size. The IPCW IBS score is reported for each method in each panel, with lower values indicating better predictive accuracy. NeuralSurv estimates the full posterior over survival functions, and the 90% credible interval is shown as a ribbon around its estimate.

Theorems & Definitions (18)

  • Theorem 3.1: Data Augmentation
  • Definition B.1: Pólya–Gamma Distribution
  • Proposition B.2
  • Proposition B.3
  • Proposition B.4
  • Theorem B.5
  • Corollary B.6
  • Definition C.1: Poisson Process
  • Definition C.2: Marked Poisson Process
  • Theorem C.3: Campbell's Theorem
  • ...and 8 more