Explainable AI for Chronic Kidney Disease Prediction Using Simulated Federated Learning

Technical Summary: Explainable AI for Chronic Kidney Disease Prediction Using Simulated Federated Learning

Problem Statement
Chronic Kidney Disease (CKD) represents a significant public health challenge, characterized by the gradual loss of kidney function. Early detection is critical for preventing severe complications and improving patient outcomes. However, developing reliable prediction models often requires aggregating sensitive patient data from multiple healthcare institutions, which raises substantial privacy concerns. Traditional centralized machine learning approaches necessitate sharing raw data with third parties, creating barriers to collaboration and data security. Furthermore, while high-accuracy models exist, the "black box" nature of many advanced algorithms hinders trust and adoption in clinical settings where transparency is paramount.

Methodology
This study proposes a framework that integrates Federated Learning (FL) with ensemble machine learning methods and Explainable AI (XAI) to predict CKD while preserving data privacy. The methodology proceeds through the following stages:

1. Dataset and Preprocessing: The research utilized a clinical dataset of 400 patient records containing 14 medical predictor variables (e.g., blood pressure, hemoglobin, serum creatinine) and a binary class label. To address class imbalance (62.5% CKD vs. 37.5% non-CKD), the Synthetic Minority Over-sampling Technique (SMOTE) was applied, expanding the dataset to 500 samples.
2. Federated Learning Architecture: The study simulated a federated environment with three distinct clients. The dataset was partitioned into three subsets using stratified splitting to preserve class distribution.
   • Local Training: Each client trained local models using three ensemble algorithms: Random Forest (RF), AdaBoost, and XGBoost.
   • Optimization: GridSearchCV was employed on the client side to automate hyperparameter tuning and prevent overfitting.
   • Model Selection: For each client, the algorithm yielding the highest accuracy on the local test set was selected as the optimal local model.
3. Global Aggregation: A central server aggregated the selected local models using a weighted VotingClassifier. The weights assigned to each model were proportional to their respective local test accuracies, ensuring that higher-performing models exerted greater influence on the global prediction.
4. Explainability: To enhance transparency, Local Interpretable Model-agnostic Explanations (LIME) were applied to interpret the global model's predictions, identifying feature contributions for individual instances.
5. Evaluation: The system was evaluated using Accuracy, Recall, Precision, and F1-Score, alongside 5-fold cross-validation to assess generalization.

Key Contributions

• Privacy-Preserving Ensemble Learning: The study demonstrates a simulated FL framework where RF, AdaBoost, and XGBoost are combined without sharing raw patient data, addressing privacy constraints in healthcare.
• Hybrid Model Optimization: By utilizing GridSearchCV for local parameter optimization and a weighted VotingClassifier for global aggregation, the approach leverages the strengths of multiple ensemble methods.
• Integration of XAI: The incorporation of LIME provides interpretability, allowing stakeholders to understand the rationale behind CKD predictions, specifically highlighting the contribution of features like hemoglobin and sugar levels.

Results

• Local Performance: In the simulated clients, Random Forest achieved 100% accuracy across all metrics for Client 1. For Clients 2 and 3, the best-performing models (AdaBoost and XGBoost, respectively) achieved approximately 97% accuracy, 100% recall, 94% precision, and 97% F1-scores.
• Global Model Performance: The aggregated global model achieved an average accuracy of 99%. Specifically, for Class 0 (non-CKD), the model achieved 98% precision and 100% recall. For Class 1 (CKD), it achieved 100% precision and 98% recall. The F1-score for both classes was 99%.
• Cross-Validation: 5-fold cross-validation confirmed the model's robustness, with accuracy scores ranging from 96% to 100% across folds, with Folds 3 and 5 reaching 100%.
• Feature Importance: LIME analysis revealed that hemoglobin levels (specifically the interval 11.30 < Hemo ≤ 13.23) and hypertension status were significant contributors to the model's decision-making process.

Significance and Claims
The paper claims that the proposed framework successfully balances high predictive performance with data privacy and model interpretability. By achieving 99% accuracy, the study highlights the potential of interpretable Federated Learning models to support early CKD diagnosis without compromising patient confidentiality. The authors assert that this approach advances data-driven healthcare solutions by offering a reliable, transparent, and privacy-preserving alternative to centralized data collection. The study concludes that such models can reduce dependency on extensive clinical testing for early detection, though it notes that future work is required to validate these findings on larger, more diverse datasets to ensure broader clinical applicability.