XGBoost-NAS with SHAP Explainability for CKD Stage Prediction: A Single-Model Federated Approach

Vrunal Sandesh Gharat, Apurv Prabhakar Patil

Abstract: Chronic Kidney Disease (CKD) is a progressive and often asymptomatic condition affecting over 850 million individuals worldwide, where early and reliable stage identification is critical for timely clinical intervention. While machine learning approaches have shown strong performance on CKD staging tasks, their practical deployment remains constrained by three key challenges: dependence on optimization patient data, poorly calibrated probability outputs, and limited interpretability for clinical decision-making. In this study, we propose a lightweight, interpretable, and privacy-preserving machine learning pipeline for multi-class CKD stage prediction based on a Bayesian-optimised XGBoost classifier deployed in a two-silo federated setting. Hyperparameter optimization is performed using Optuna-based Neural Architecture Search (NAS), while model outputs are calibrated using Platt scaling to ensure reliable probabilistic predictions. Model decisions are further explained using SHAP-based global and local interpretability techniques. The proposed approach is evaluated on a dataset of 15,736 longitudinal CKD records across all five KDIGO stages and benchmarked against five baseline models, including Logistic Regression, Random Forest, SVM, standard XGBoost, and a Multi-Layer Perceptron. The federated NAS-XGBoost model achieves an AUROC of 0.9748, AUPRC of 0.9512, weighted F1-score of 0.8863, and Expected Calibration Error (ECE) of 0.0312, outperforming all baselines in both discrimination and calibration. To address the potential dependency on eGFR, an ablation study is conducted by excluding this primary clinical feature, demonstrating that while eGFR contributes the dominant signal, secondary biomarkers retain meaningful predictive capacity. This highlights the model?s robustness in handling realistic clinical variability. The results indicate that a single, well-calibrated, and interpretable model can achieve competitive performance in a federated setting without requiring data optimization, supporting its applicability in privacy-sensitive healthcare environments.

Keywords: Chronic Kidney Disease, CKD Stage Prediction, XGBoost, Neural Architecture Search, Federated Learning, SHAP, Platt Scaling, Calibration, eGFR, KDIGO, Explainable AI

How to Cite?: Vrunal Sandesh Gharat, Apurv Prabhakar Patil, "XGBoost-NAS with SHAP Explainability for CKD Stage Prediction: A Single-Model Federated Approach", Volume 15 Issue 6, June 2026, International Journal of Science and Research (IJSR), Pages: 945-952, https://www.ijsr.net/getabstract.php?paperid=SR26617223814, DOI: https://dx.doi.org/10.21275/SR26617223814

Download Citation: APA | MLA | BibTeX | EndNote | RefMan