Efficient Training of Probabilistic Neural Networks for Survival Analysis

Area of Science:

• Machine Learning
• Biostatistics
• Computational Biology

Background:

• Variational Inference (VI) is standard for Bayesian inference in deep learning but computationally expensive, especially in high dimensions.
• Existing alternatives like Monte Carlo Dropout (MCD) and Spectral-normalized Neural Gaussian Process (SNGP) are underutilized in survival analysis.
• Training deep probabilistic survival models on large datasets necessitates efficient inference methods with low computational overhead.

Purpose of the Study:

• To investigate non-Variational Inference (VI) methods for training deep probabilistic survival models in large, high-dimensional datasets.
• To evaluate the prediction performance, uncertainty calibration, and model complexity of VI, MCD, and SNGP in survival analysis.
• To determine if MCD and SNGP can match or exceed VI's performance without increasing model complexity.

Main Methods:

• Implemented and compared three probabilistic approaches: Variational Inference (VI), Monte Carlo Dropout (MCD), and Spectral-normalized Neural Gaussian Process (SNGP).
• Evaluated methods on the MIMIC-IV dataset, assessing prediction accuracy (concordance index, mean absolute error) and uncertainty calibration.
• Analyzed model complexity and computational overhead associated with each inference technique.

Main Results:

• MCD demonstrated comparable prediction performance to VI on the MIMIC-IV dataset (concordance index: 0.748 vs. 0.743; MAE: 254.9 vs. 254.7).
• MCD provided C-calibrated uncertainty estimates, while SNGP achieved D-calibrated survival functions across all tested datasets (4/4), outperforming VI (2/4).
• Non-VI methods showed potential for comparable or improved performance with reduced computational complexity.

Conclusions:

• Monte Carlo Dropout (MCD) and Spectral-normalized Neural Gaussian Process (SNGP) are viable, computationally efficient alternatives to Variational Inference (VI) for deep probabilistic survival analysis.
• These methods offer competitive prediction accuracy and improved uncertainty calibration, particularly in high-dimensional settings.
• The findings encourage wider adoption of MCD and SNGP for survival analysis in large datasets where computational efficiency is critical.