Predicting arrhythmia recurrence post-ablation in atrial fibrillation using explainable machine learning

Abstract

Background: Following atrial fibrillation ablation, it is challenging to distinguish patients who will remain arrhythmia-free from those at risk for recurrence. New explainable machine learning (xML) techniques allow for systematic assessment of arrhythmia recurrence risk following catheter ablation. We aim to develop an xML algorithm that predicts recurrence and reveals key risk factors to facilitate better follow-up strategy after an ablation procedure.

Methods: We reconstructed pre-and post-ablation models of the left atrium (LA) from late gadolinium enhanced magnetic resonance (LGE-MRI) for 67 patients. Patient-specific features (LGE-based measurements of pre/post-ablation arrhythmogenic substrate, LA geometry metrics, computational simulation results, and clinical risk factors) trained a random forest classifier to predict recurrent arrhythmia. We calculated each risk factor's marginal contribution to model decision making via SHapley Additive exPlanations (SHAP).

Results: The classifier accurately predicts post-ablation arrhythmia recurrence (mean receiver operating characteristic [ROC] area under the curve [AUC]: 0.80 ± 0.04; mean precision-recall [PR] AUC: 0.82 ± 0.08). SHAP analysis reveals that of 89 features tested, the key population risk factors for recurrence are: large left atrium, low LGE-quantified post-ablation scar in the atrial floor region, and previous attempts at direct current cardioversion. We also examine patient-specific recurrence predictions, since xML allows us to understand why a particular individual can have large prediction weights for some categories without tipping the balance towards an incorrect prediction. Finally, we validate our model in a completely new, 15-patient retrospective holdout cohort (80% correct).

Conclusion: Our SHAP-based explainable machine learning approach is a proof-of-concept clinical tool to explain arrhythmia recurrence risk in patients who underwent ablation by combining patient-specific clinical profiles and LGE-derived data.

Fig. 1. Flowchart outlining the complete model development process.

A To address multicollinearity between the 89 risk factors compiled for this study, LASSO regression removed the least important and collinear features. B The product of this first stage was the 27-element LASSO-optimized feature set (LOFS). C Subsequently, the LOFS were used to train and test either random forest machine learning (ML) or logistic regression models using five-fold, 80:20-split cross-validation. The random forest and logistic regression models were then tested on data from a never-before-seen 15-patient holdout cohort. To assess model explainability, the marginal contributions of individual LOFS values on overall random forest model predictions in the original and holdout cohorts were evaluated by SHAP analysis. SHAP analysis was not needed for the logistic regression model, as each feature had a coefficient explicitly describing its impact on model predictions. Holdout and explainability tests were always performed on the single best logistic regression or random forest model, as assessed during the predictive efficacy stage via the area under the receiver operating characteristic curve (AUROC) metric. Links to relevant figures later in the study in which specific results are presented are provided in Output panels.

Fig. 2. Random forest prediction algorithm efficacy in testing sets across five folds.

A ROC curves for testing set (AUC: 0.80 ± 0.04). B Precision-Recall curves for testing set. (AUC: 0.82 ± 0.08).

Fig. 3. Summary of feature importance.

Features are sorted in descending order of importance. Features are classified as either LA geometry (black), clinical attributes (blue), fibrosis (green), ablation-delivered scar (orange), or results from biophysically detailed simulations (magenta). Red dots indicate patients with a high feature value for continuous variables, or presence of variable in binary cases. The vertical center line represents no impact on model outcome, while right-shifted data suggests association with recurrence and left-shifted data indicates non-recurrence. AAD antiarrhythmic drug, RPV right pulmonary veins.

Fig. 4. Dependence plots of SHAP importance vs. raw feature values.

SHAP values greater than zero correlate to increased risk of recurrence, and those less than 0 correlate with reduced recurrence risk. A Dependence plot for post-ablation LA volume index (LAVI), demonstrating increased recurrence in patients with a post-ablation LAVI above 51.121 mL/m². y = 0.220/(1+e^(−0.668*(x-51.121))–0.107. B Dependence plot for post-ablation scar in the atrial floor region, showing less scar on the atrial floor was associated with increased recurrence. y = 0.0934*e^(–0.788x)–0.031. C Dependence plot for number of direct current (DC) cardioversions, indicating recurrence was more likely in patients who received more DC cardioversions. y = 0.0306x–0.0366; R² = 0.8699. D Dependence plot for post-ablation scar in the region of the RPVs, fit with a 7-period moving average trendline, suggesting the random forest model identified a complex relationship between scar in the region of the RPVs and recurrence risk.

Fig. 5. Explanation of individual patient prediction scores for representative patients in the original dataset.

Yellow and magenta arrows indicate the marginal contribution of each feature, and the grey outline indicates the model’s prediction. Black indicator lines provide more detailed information about the features corresponding to the arrows they intersect. A Visualization of pre- and post-ablation models and explanation of a correct recurrent arrhythmia prediction driven by geometry and substrate features. This patient recurred with atrial fibrillation (AFib). B Pre- and post-ablation models of patient and explanation of a correct non-recurrent arrhythmia prediction driven by geometry and clinical features. C Visualization of pre- and post-ablation patient-specific models and explanation of a correct non-recurrent arrhythmia prediction driven by pre- and post-ablation substrate features. This patient also recurred with AFib.

Fig. 6. Summary of model performance on a previously unseen internal validation (holdout) cohort.

A Chart of each patient outcome and model prediction in the 15-patient cohort. Gradient indicates relative confidence (summed SHAP values offset by the model’s decision threshold of 0.10) of non-recurrence and recurrence, respectively, calculated as the sum of SHAP values. B Example for correct prediction with post-ablation patient-specific anatomical model and prediction breakdown via SHAP analysis. C Example for incorrect positive prediction.

Fig. 7. Performance of a model trained with the change in LA volume index (LAVI) in place of the pre-ablation LAVI.

A ROC curve from five-fold training and cross-validation of the model, showing a mean ROC AUC of 0.85 ± 0.03, with individual fold ROC AUCs ranging from 0.80 to 0.90. The model retained its 80% accuracy on the holdout set. B Dependence plot showing the relationship between the change in LAVI and the marginal contribution to model outcome.