Interpretable machine learning model for new-onset atrial fibrillation prediction in critically ill patients: a multi-center study

Abstract

Background: New-onset atrial fibrillation (NOAF) is the most common arrhythmia in critically ill patients admitted to intensive care and is associated with poor prognosis and disease burden. Identifying high-risk individuals early is crucial. This study aims to create and validate a NOAF prediction model for critically ill patients using machine learning (ML).

Methods: The data came from two non-overlapping datasets from the Medical Information Mart for Intensive Care (MIMIC), with MIMIC-IV used for training and subset of MIMIC-III used as external validation. LASSO regression was used for feature selection. Eight ML algorithms were employed to construct the prediction model. Model performance was evaluated based on identification, calibration, and clinical application. The SHapley Additive exPlanations (SHAP) method was used for visualizing model characteristics and individual case predictions.

Results: Among 16,528 MIMIC-IV patients, 1520 (9.2%) developed AF post-ICU admission. A model with 23 variables was built, with XGBoost performing best, achieving an AUC of 0.891 (0.873-0.888) in validation and 0.769 (0.756-0.782) in external validation. Key predictors included age, mechanical ventilation, urine output, sepsis, blood urea nitrogen, percutaneous arterial oxygen saturation, continuous renal replacement therapy and weight. A risk probability greater than 0.6 was defined as high risk. A friendly user interface had been developed for clinician use.

Conclusion: We developed a ML model to predict the risk of NOAF in critically ill patients without cardiac surgery and validated its potential as a clinically reliable tool. SHAP improves the interpretability of the model, enables clinicians to better understand the causes of NOAF, helps clinicians to prevent it in advance and improves patient outcomes.

Keywords: Critically ill patients; MIMIC database; Machine learning; New-onset atrial fibrillation; Predictive models.

Fig. 1

Patient screening flow from the MIMIC database. NOAF: new-onset atrial fibrillation

Fig. 2

Lasso regression-based variable screening. A. Variation characteristics of variable coefficients; B. The process of selecting the optimal value of the parameter λ in the lasso regression model is carried out by the cross-validation method

Fig. 3

ROC curves for the machine learning models. XGBoost: extreme gradient boosting; SVM: support vector machine; Adaboost: adaptive boosting; MLP: multilayer perceptron; NN: neural network; NB: naive bayes; LR: logistic regression; GBM: gradient boosting machine; ROC: receiver operating characteristic; AUC: area under the curve

Fig. 4

Calibration capability and clinical benefit of the model. A. Calibration curve B. Clinical Impact Curve (CIC) C. Decision curve analysis (DCA), XGBoost: extreme gradient boosting; SVM: support vector machine; Adaboost: adaptive boosting; MLP: multilayer perceptron; NN: neural network; NB: naive bayes; LR: logistic regression; GBM: gradient boosting machine.

Fig. 5

Visually interpret machine learning models using SHAP. A SHAP summary point. B SHAP force plot. SBP: systolic blood pressure; BUN: blood urea nitrogen; SpO₂: percutaneous arterial oxygen saturation; WBC: white blood cell; MBP: mean blood pressure; DBP: diastolic blood pressure; HFrEF: heart failure with reduced ejection fraction; HFpEF: heart failure with preserved ejection fraction