Source-linked AI summary
A Scalable Discrete-Time Survival Model for Neural Networks
Michael F. Gensheimer, Balasubramanian Narasimhan
TL;DR
Neural-network prediction in medicine needs methods that use censored survival data without discarding follow-up information. The paper presents Nnet-survival, a flexible discrete-time model trained by maximum likelihood with minibatch SGD, and reports good performance on simulated and real data alongside practical advantages for larger datasets.
Problem
Medical prediction models often evaluate outcomes at one follow-up time and provide limited treatment of censored data, losing information and requiring retraining for different time points.
Method
Nnet-survival is a discrete-time survival model trained with a likelihood-based loss using minibatch gradient descent, allowing baseline hazards and predictor effects to vary with follow-up time.
Results
Nnet-survival produced good discrimination and calibration on simulated and real data, while running time was similar to Deepsurv and better than Cox-nnet for sample sizes >1000.
Takeaways & Limitations
The model may be useful for non-proportional hazards, rapid training, large datasets, and neural-network inputs such as images or text.
Takeaways & Limitations
Discretizing follow-up time produces less smooth survival curves, and the model does not provide predictions past the last time interval.
Abstract
from arXiv · showhide
There is currently great interest in applying neural networks to prediction tasks in medicine. It is important for predictive models to be able to use survival data, where each patient has a known follow-up time and event/censoring indicator. This avoids information loss when training the model and enables generation of predicted survival curves. In this paper, we describe a discrete-time survival model that is designed to be used with neural networks, which we refer to as Nnet-survival. The model is trained with the maximum likelihood method using minibatch stochastic gradient descent (SGD). The use of SGD enables rapid convergence and application to large datasets that do not fit in memory. The model is flexible, so that the baseline hazard rate and the effect of the input data on hazard probability can vary with follow-up time. It has been implemented in the Keras deep learning framework, and source code for the model and several examples is available online. We demonstrate the performance of the model on both simulated and real data and compare it to existing models Cox-nnet and Deepsurv.
INTRODUCTION
Neural-network outcome models often use a single follow-up point and inadequately address censoring, risking lost patient information and retraining for different time points. Nnet-survival is proposed as a theoretically justified discrete-time model that handles non-proportional hazards and supports rapid minibatch training.
- Motivation: Single-time-point outcome models often do not adequately address censored patients, losing information about their follow-up.Such models would also need retraining to predict outcomes at different time points.
- Prior approaches: Existing neural-network survival approaches commonly adapt linear or Cox models and often assume proportional hazards.The cited alternatives include heuristic-loss discrete-time, Weibull parametric, and Cox-model adaptations.
- Motivation: Modern datasets with thousands or millions of patients can expose violations of the proportional hazards assumption through residual plots or statistical tests.
- Proposed model: Nnet-survival is a theoretically justified discrete-time survival model designed to handle non-proportional hazards and train rapidly with minibatch gradient descent.
- Intended use: The model is positioned for known non-proportional hazards, very large datasets that do not fit in memory, and neural-network-suitable predictors such as images or text.
MATERIALS & METHODS
Nnet-survival is a discrete-time neural-network survival model built around a likelihood-based loss, minibatch SGD, and flexible or proportional-hazards parameterizations. It represents follow-up with interval-specific survival probabilities and accommodates censored observations during training.
- Model formulation: Follow-up time is divided into fixed intervals, with a conditional hazard probability for failure in each interval given survival to its beginning.The hazard may vary by individual according to neural-network inputs and weights.
- Likelihood and censoring: The loss uses failure and at-risk counts within each interval, while censored individuals contribute survival information through their censoring interval.Individuals censored late in an interval receive credit for surviving that interval to avoid downward bias in survival estimates.
- Model formulation: The proposed negative log-likelihood loss is theoretically grounded in established discrete-time survival models rather than relying on a heuristic loss.The likelihood formulation can be optimized by stochastic or minibatch gradient descent.
- Hazard parameterizations: The flexible version produces a separate hazard rate for each time interval, allowing both baseline hazard and predictor effects to vary over follow-up time.Its n-dimensional output has one node per interval and is fully connected to the last hidden layer.
- Hazard parameterizations: The proportional-hazards version allows baseline hazard to vary by interval while constraining predictor effects to remain constant over follow-up time.This version reduces parameters and is useful for smaller datasets or interpretability.
- Implementation: Nnet-survival was implemented in Python with Keras and TensorFlow, using an n-dimensional output whose elements are conditional interval-survival probabilities.Cumulative survival through interval j is computed from the interval-specific predictions.
RESULTS
Nnet-survival showed strong calibration on simulated data, MNIST-derived survival prediction, and the SUPPORT study, while maintaining reliable convergence and robustness to interval choices. In large-data timing tests, Cox-nnet exceeded available memory at sample sizes of 100,000 or more.
- Simulated data: Actual and predicted survival curves were superimposed for the two simulated predictor groups, indicating excellent calibration.The groups had median survival times of 200 and 400 days.
- Convergence: Repeated optimization from different random starting weights converged to very similar final loss or likelihood values.This indicates reliable model convergence across the tested starting conditions.
- Time intervals: Discrimination was identical at C-index 0.66 across four time-interval widths and configurations, showing robustness to interval choices.The evaluated configurations used constant width or increasing width with increasing follow-up time.
- MNIST dataset: Calibration was excellent on MNIST, with a test-set C-index of 0.713 versus 0.770 for a perfect model using the true digit.Digits 0 through 4 represented predictors of survival time, with lower digits having longer median survival.
- SUPPORT study: Several SUPPORT predictors violated proportional hazards, including metastatic cancer with similar early-death risk but higher late-death risk.The violation was evidenced by non-parallel lines in the log(-log survival) plot.
- SUPPORT study: On the SUPPORT test set, Nnet-survival had the best calibration at all three follow-up times, while discrimination was similar across models with C-index around 0.73.Visual calibration was strongest for Nnet-survival at 6 months and 1 year; Cox-nnet and standard Cox tended to under-predict survival for the best-prognosis patients.
- Running time: Cox-nnet ran out of memory at sample sizes of 100,000 and higher on a computer with 32 GB of memory.Its running time could therefore only be calculated for sample sizes from 1,000 to 31,622.
DISCUSSION
Nnet-survival combines likelihood-based discrete-time modeling with minibatch training, supporting non-proportional hazards and large datasets. It performed comparably to established survival models on SUPPORT data, while retaining practical limitations from time discretization and finite prediction horizons.
- DISCUSSION: Nnet-survival uses the likelihood function as its loss, naturally incorporates non-proportional hazards, and supports multiple neural-network architectures.The model can be applied to multilayer perceptrons and convolutional neural networks.
- DISCUSSION: On the SUPPORT study dataset of 9,105 hospitalized patients, discrimination and calibration were similar to several neural-network survival models and traditional Cox regression.The dataset had moderate sample size and low-dimensional predictors that might not benefit substantially from neural networks.
- DISCUSSION: For sample sizes >1000, Nnet-survival ran in a time similar to Deepsurv and better than Cox-nnet.Cox-nnet ran out of memory for larger datasets because it stores an n by n matrix.
- DISCUSSION: Discretizing follow-up time produces less smooth survival curves than non-discrete parametric models, although sufficient intervals can make curves appear very smooth.The authors note that curves using 19 intervals appeared very smooth.
- DISCUSSION: Unlike parametric survival models, Nnet-survival does not provide predictions beyond the final time interval.The authors recommend extending the last interval past the last follow-up time of interest.
- DISCUSSION: Future flexible parametric models could combine parametric-survival advantages with complex non-proportional hazards and deep-learning implementations.The paper identifies cubic spline-based models as one possible direction.
CONCLUSIONS
The model supports non-proportional hazards and stochastic-gradient training, and showed good discrimination and calibration on simulated and real data.
- CONCLUSIONS: Nnet-survival allows non-proportional hazards, stochastic-gradient training, rapid training, and good discrimination and calibration on simulated and real data.