Estimating survival rates using artificial intelligence combined with the Aalen–Johansen estimator in multi-state models

Abstract

Accurate survival prediction is essential for clinical decision-making, health economics, and treatment planning. Traditional methods like the Kaplan-Meier and Cox models are widely used but have limitations when applied to complex multi-state processes or individualized predictions. The Aalen–Johansen estimator, a non-parametric approach suited for multi-state Markov models, improves population-level inference but lacks the ability to incorporate covariates or capture nonlinear relationships. In this study, we propose a hybrid framework that combines the Aalen–Johansen estimator with artificial intelligence (AI) techniques, specifically gradient boosting machines (GBM) and long short-term memory (LSTM) networks. By transforming transition probabilities into subject-level pseudo-observations, AI models can learn personalized survival functions based on individual covariates. We validate our approach on both simulated and real-world clinical datasets. The hybrid model outperforms traditional estimators in predictive accuracy, as measured by calibration and discrimination metrics such as Brier score and area under the curve (AUC). This AI–Aalen–Johansen framework enhances risk stratification and clinical decision-making by providing more accurate, scalable, and interpretable survival predictions. Our results support its potential as a valuable tool in modern healthcare analytics, contributing to the advancement of precision medicine.