Neural Diffusion Processes for Physically Interpretable Survival Prediction

TL;DR

This work tackles non-proportional survival prediction by introducing DeepFHT, a neural framework that couples deep learning with first hitting time (FHT) distributions of latent diffusion processes to an absorbing boundary. A neural network maps covariates to physically meaningful diffusion parameters (initial condition, drift, diffusion), enabling closed-form survival $S(t)$ and hazard $h(t)$ functions under Brownian and arithmetic Brownian motion. The model achieves competitive predictive accuracy on public clinical datasets and offers physics-based interpretability by locating individuals in a latent parameter space where distances reflect risk similarity and feature-risk associations. Overall, DeepFHT provides a principled, interpretable alternative to black-box survival models and demonstrates the value of integrating stochastic-process theory with deep learning for complex time-to-event phenomena.

Abstract

We introduce DeepFHT, a survival-analysis framework that couples deep neural networks with first hitting time (FHT) distributions from stochastic process theory. Time to event is represented as the first passage of a latent diffusion process to an absorbing boundary. A neural network maps input variables to physically meaningful parameters including initial condition, drift, and diffusion, within a chosen FHT process such as Brownian motion, both with drift and driftless. This yields closed- form survival and hazard functions and captures time-varying risk without assuming proportional- hazards. We compare DeepFHT with Cox regression using synthetic and real-world datasets. The method achieves predictive accuracy on par with the state-of-the-art approach, while maintaining a physics- based interpretable parameterization that elucidates the relation between input features and risk. This combination of stochastic process theory and deep learning provides a principled avenue for modeling survival phenomena in complex systems

Figures (5)

• Figure 1: Example output of the Lévy FHT model. Individual-specific survival functions are computed from the neural network–predicted parameters.
• Figure 2: Performance across clinical and synthetic datasets. Scatterplots with error bars for C-index (↑, brown, top row) and IBS (↓, dark blue, bottom row) across datasets (GBSG2, Framingham, SUPPORT2, NonPH). For each dataset, three models are shown: Lévy, inverse Gaussian, and CoxPH. Mean values and standard deviations are computed via 100 bootstrap iterations.
• Figure 3: Event times in the parameter spaces of Deep FHT models. Left: Framingham dataset in the space of Lévy Deep FHT ($\{x_0,D\}$). Right: Framingham dataset in the space of inverse Gaussian Deep FHT ($\{x_0,\mu\}$). Background colors represent interpolated event times $T(\mathbf{p})$ obtained by inverse distance weighting of uncensored training instances. Points correspond to uncensored test patients and individuals surviving beyond the last uncensored time, colored by their observed event times. In both cases, patients with similar event times cluster in contiguous regions, illustrating that the model encodes similarity in terms of process parameters.
• Figure 4: Feature–parameter relationships in the Lévy and inverse Gaussian DeepFHT models. Top: Framingham dataset with Lévy model, showing the parameter space $\{x_0,D\}$ colored by systolic (left) and diastolic (right) blood pressure. Bottom: GBSG2 dataset with inverse Gaussian model, showing the parameter space $\{x_0,\mu\}$ for patients with grade 1 (left) and grade 3 (right) tumors. In both cases, clinical risk factors align with model-derived high-risk regions (small $x_0$, large $D$ or $\mu$), supporting the physical interpretability of the parameterization.
• Figure 5: Time interpolation in parameter space across models for GBSG2, SUPPORT2 and NonPH datasets. Notice the absence of collinearity between parameters for the synthetic NonPH datasets, as discussed in Sec. \ref{['sec:PX']}