基于Transformer算法构建的缺血性脑卒中早期预测模型

An Early Prediction Model for Ischemic Stroke Based on the Transformer Algorithm

摘要: 目的：采用Transformer等7种机器学习算法构建缺血性脑卒中风险预测模型，筛选最优模型并进行可解释性分析，助力早期识别高风险人群。方法：采用Kaggle公开脑卒中数据集，纳入5110例患者，其中病发组249例、非病发组4860例。经填补缺失值、SMOTE平衡样本后，分别采用BP、RF、XGBoost、SVM、DNN、MLP、Transformer算法构建预测模型。通过五折交叉验证进行内部验证，以AUC、准确率、召回率、特异度、Brier分数评价模型性能，并采用SHAP算法对最优模型进行可视化解释。结果：Transformer模型的AUC为0.912、准确率0.771、召回率0.847、特异度0.766、Brier分数0.091，综合性能优于其他6种模型。SHAP分析显示，年龄、工作类型、BMI、心脏病是影响模型预测贡献度最高的特征。结论：基于Transformer算法构建的缺血性脑卒中早期预测模型具有较好的预测效能，结合SHAP可解释性分析可辅助临床识别高危因素，为Transformer算法在卒中流行病学分析中的适用性提供了初步依据。

Abstract: Objective: This study aimed to develop and compare seven machine learning algorithms, including Transformer, to build risk prediction models for ischemic stroke, identify the best-performing model, and perform interpretability analysis to support early detection of high-risk individuals. Methods: A publicly available Kaggle stroke dataset was used, comprising 5110 patients (249 with stroke, 4860 without). Missing values were imputed using random forests, and the dataset was balanced with SMOTE. Prediction models were built using BP, RF, XGBoost, SVM, DNN, MLP, and Transformer algorithms. Internal validation was conducted via five-fold cross-validation. Model performance was evaluated using AUC, accuracy, recall, specificity, and Brier score. SHAP was applied to interpret the best model visually. Results: The Transformer model achieved an AUC of 0.912, accuracy of 0.771, recall of 0.847, specificity of 0.766, and Brier score of 0.091, outperforming the other six models. SHAP analysis revealed that age, work type, BMI, and heart disease were the most influential features. Conclusion: The Transformer-based early prediction model for ischemic stroke shows strong predictive performance. Combined with SHAP for interpretability, it can help clinicians identify key risk factors and support early prevention strategies.