Early Prediction of Stroke Risk: A Machine Learning Practice Based on Ensemble Learning and SHAP Interpretability

This project stems from the developer's family experience and aims to build an explainable early stroke risk prediction tool. Core technologies include ensemble learning (XGBoost, Random Forest, etc.), SMOTE for handling class imbalance issues, and SHAP interpretability analysis to provide reliable risk assessment support for medical screening. The project code is open-source (GitHub link: https://github.com/viscl/stroke-risk) and was released in 2026.

Stroke is the second leading cause of death globally and the primary cause of adult disability. This project was born from the developer's personal family experience with stroke, with the core goal of creating an easily accessible and highly interpretable screening tool to identify high-risk groups in advance and assist in preventive interventions.

The Kaggle Stroke Prediction Dataset was used (5110 records, 5% of which are stroke patients), containing 10 features:

• Demographics: Gender, Age, Marital Status, Residence Type
• Health Indicators: Hypertension, Heart Disease, Average Blood Glucose Level, BMI (missing values filled with median)
• Lifestyle: Work Type, Smoking Status The target variable is whether a stroke occurred (binary classification).

The technical process includes:

1. Preprocessing: OneHot encoding for categorical features, standardization for numerical features, missing value handling
2. Model Selection: Integrating XGBoost (gradient boosting), Random Forest (Bagging), Logistic Regression (baseline), Neural Network (non-linear interaction)
3. Class Imbalance Handling: SMOTE technology (applied only to cross-validation training sets to avoid data leakage)
4. Evaluation: 5-fold stratified cross-validation to ensure generalization performance.

SHAP is introduced to solve the black-box problem of medical AI; its values are additive, consistent, and fair. TreeExplainer is used to interpret tree models. Risk grading: <30% low risk, 30%-60% medium risk, >60% high risk. Example: A 67-year-old male (with heart disease history, high blood glucose, obesity, and former smoking) was predicted as high risk, and SHAP identified age and blood glucose as the main driving factors.

Value: Priority on interpretability (SHAP), correct handling of class imbalance, multi-model comparison, clinically friendly risk grading. Limitations: Small dataset size (5110 records), extreme class imbalance, incomplete feature coverage (lack of family history, etc.), regional generalization needs verification, should be used as an auxiliary tool for doctors (not a replacement).

This project demonstrates the practice of transforming machine learning into a medical tool, emphasizing interpretability, class imbalance handling, and practical application orientation. It provides a learning case for medical AI developers, helping with stroke prevention and the popularization of medical AI.