[Paper Review] Neural ODEs for Multi-State Survival Analysis.

Stefan Groha, Sebastian M. Schmon+1 more|arXiv (Cornell University)|Jun 8, 2020

Machine Learning in Healthcare4 references4 citations

TL;DR

This paper proposes neural ordinary differential equations (Neural ODEs) to model multi-state survival processes by solving the Kolmogorov forward equations, enabling flexible, non-linear hazard rate estimation with individual-specific trajectories. It further integrates a variational latent variable model to quantify uncertainty in cause-specific hazard rates, achieving state-of-the-art performance on benchmark survival datasets.

ABSTRACT

Survival models are a popular tool for the analysis of time to event data with applications in medicine, engineering, economics and many more. Advances like the Cox proportional hazard model have enabled researchers to better describe hazard rates for the occurrence of single fatal events, but are limited by modeling assumptions, like proportionality of hazard rates and linear effects. Moreover, common phenomena are often better described through multiple states, for example, the progress of a disease might be modeled as healthy, sick and dead instead of healthy and dead, where the competing nature of death and disease has to be taken into account. Also, individual characteristics can vary significantly between observational units, like patients, resulting in idiosyncratic hazard rates and different disease trajectories. These considerations require flexible modeling assumptions. Current standard models, however, are often ill-suited for such an analysis. To overcome these issues, we propose the use of neural ordinary differential equations as a flexible and general method for estimating multi-state survival models by directly solving the Kolmogorov forward equations. To quantify the uncertainty in the resulting individual cause-specific hazard rates, we further introduce a variational latent variable model. We show that our model exhibits state-of-the-art performance on popular survival data sets and demonstrate its efficacy in a multi-state setting.

Motivation & Objective

• To address limitations in traditional survival models, such as proportional hazards and linear effects, which fail to capture complex, non-linear disease progression.
• To model multi-state processes—such as healthy → sick → dead—where competing risks and state transitions are critical for accurate analysis.
• To allow individual-specific hazard rate estimation by capturing idiosyncratic disease trajectories across observational units like patients.
• To introduce a variational latent variable model to quantify uncertainty in estimated cause-specific hazard rates.
• To develop a general, flexible framework that directly solves the Kolmogorov forward equations using Neural ODEs for improved modeling of dynamic survival processes.

Proposed method

• Utilizes neural ordinary differential equations (Neural ODEs) to parameterize the transition intensity functions in multi-state survival models.
• Solves the Kolmogorov forward equations numerically via Neural ODEs to model the time evolution of state probabilities.
• Introduces a variational latent variable model to estimate uncertainty in individual-specific cause-specific hazard rates.
• Employs a continuous normalizing flow-like structure to model the latent space, enabling flexible and differentiable density estimation.
• Trains the model end-to-end using likelihood-based optimization, allowing joint estimation of transition dynamics and uncertainty.
• Models individual heterogeneity by learning patient-specific latent representations that influence hazard rates through the Neural ODE architecture.

Experimental results

Research questions

• RQ1 Can Neural ODEs effectively model complex, non-linear hazard rate dynamics in multi-state survival processes compared to standard parametric models?
• RQ2 How well can a variational latent variable model quantify uncertainty in individual cause-specific hazard rates within a multi-state framework?
• RQ3 To what extent does the proposed method outperform existing survival models on real-world multi-state survival datasets?
• RQ4 Can the model capture individual-specific disease trajectories and competing risks without assuming proportional hazards or linear effects?
• RQ5 How robust is the method to model misspecification and varying data structures in multi-state survival analysis?

Key findings

• The proposed Neural ODE-based model achieves state-of-the-art performance on popular survival datasets, outperforming standard models in predictive accuracy.
• The integration of a variational latent variable model enables reliable quantification of uncertainty in individual cause-specific hazard rates.
• The model successfully captures non-linear, time-varying hazard rates and individual-specific trajectories without assuming proportional hazards.
• By solving the Kolmogorov forward equations via Neural ODEs, the method provides a continuous and differentiable framework for modeling state transitions.
• The approach effectively handles competing risks in multi-state processes, such as disease progression and death, by modeling transitions between multiple states simultaneously.
• Empirical results demonstrate the model's flexibility and generalization capability across diverse survival data structures.