Survival Mixture Density Networks

TL;DR

Survival Mixture Density Networks offer an efficient, flexible alternative to neural ODE-based models for survival analysis, achieving comparable or better performance with faster training times and avoiding binning issues.

Contribution

This paper introduces Survival MDNs, a novel continuous time survival analysis model that is faster and more flexible than neural ODE-based approaches.

Findings

Survival MDNs outperform or match baseline models on key metrics.

They are significantly faster to train than neural ODE models.

Survival MDNs avoid binning issues present in discrete models.

Abstract

Survival analysis, the art of time-to-event modeling, plays an important role in clinical treatment decisions. Recently, continuous time models built from neural ODEs have been proposed for survival analysis. However, the training of neural ODEs is slow due to the high computational complexity of neural ODE solvers. Here, we propose an efficient alternative for flexible continuous time models, called Survival Mixture Density Networks (Survival MDNs). Survival MDN applies an invertible positive function to the output of Mixture Density Networks (MDNs). While MDNs produce flexible real-valued distributions, the invertible positive function maps the model into the time-domain while preserving a tractable density. Using four datasets, we show that Survival MDN performs better than, or similarly to continuous and discrete time baselines on concordance, integrated Brier score and integrated binomial log-likelihood. Meanwhile, Survival MDNs are also faster than ODE-based models and circumvent binning issues in discrete models.