Joint prediction of chronic kidney disease severity and renal function biomarkers using attention-driven multi-task learning
Chronic kidney disease (CKD) is a progressive and largely irreversible condition that often remains undetected until advanced stages, and its clinical assessment depends jointly on disease stage, glomerular filtration rate (GFR), and serum creatinine. Most published machine learning models predict only one of these targets and therefore ignore the physiological dependencies linking them. This study introduces an attention-based multi-task learning (ABMTL) framework that predicts all three simultaneously from structured clinical data. A shared representation backbone learns common patterns of renal dysfunction, while three task-specific attention modules re-weight the shared latent features for each objective and feed independent classification and regression heads trained under a joint loss. A leakage-free preprocessing protocol combined correlation-driven adaptive imputation, smoothed target encoding, and training-set-only normalization. The framework was evaluated on 13,500 patient records comprising 13 clinical features, and compared against seven machine learning and deep learning baselines under identical partitions. ABMTL achieved 95.55% accuracy for six-class CKD staging, with coefficients of determination of 0.966 for GFR and 0.976 for serum creatinine, outperforming every baseline on all metrics. Ablation experiments confirmed that the attention modules and multi-task optimization contribute complementary gains, and paired testing indicated the improvements were statistically significant. Kernel Shapley additive explanations analysis identified albumin, blood urea, hemoglobin, hypertension, and diabetes as the dominant predictors, consistent with established nephrology. Serum creatinine and GFR were used only as prediction targets and never as inputs. At the clinically relevant cut-offs of stage ≥3 and stage 4–5, sensitivity exceeded 97% and specificity exceeded 99%. Because the design is single-center, retrospective, and cross-sectional, and calibration and decision-curve analyses have not yet been performed, the framework is presented as a candidate screening-support model that requires external, prospective, and outcome-based validation before clinical use.

- Mallamaci F, Tripepi G. Risk factors of chronic kidney disease progression: between old and new concepts. J Clin Med. 2024;13(3):678. doi: 10.3390/jcm13030678
- Mark PB, Stafford LK, Gram ME, et al. Global, regional, and national burden of chronic kidney disease in adults, 1990-2023, and its attributable risk factors: a systematic analysis for the Global Burden of Disease Study 2023. Lancet. 2025;406(10518):2461-2482. doi: 10.1016/S0140-6736(25)01853-7
- Tangri N, Peach EJ, Franzén S, Barone S, Kushner PR. Patient management and clinical outcomes associated with a recorded diagnosis of stage 3 chronic kidney disease: the REVEAL-CKD study. Adv Ther. 2023;40(6):2869-2885. doi: 10.1007/s12325-023-02482-5
- Stevens PE, Ahmed SB, Carrero JJ, et al. KDIGO 2024 clinical practice guideline for the evaluation and management of chronic kidney disease. Kidney Int. 2024;105(4S):S117-S314. doi: 10.1016/j.kint.2023.10.018
- Khandker SS, Kundu S, Ahmed F, et al. Association of serum creatinine, urea, and glomerular filtration rate with the progression of diabetic associated kidney complications: a retrospective case-control study. Curr Issues Mol Biol. 2026;48(2):167. doi: 10.3390/cimb48020167
- Inker LA, Eneanya ND, Coresh J, et al. New creatinine- and cystatin C-based equations to estimate GFR without race. N Engl J Med. 2021;385(19):1737-1749. doi: 10.1056/NEJMoa2102953
- Delgado C, Baweja M, Crews DC, et al. A unifying approach for GFR estimation: recommendations of the NKF-ASN Task Force on Reassessing the Inclusion of Race in Diagnosing Kidney Disease. J Am Soc Nephrol. 2021;32(12):2994-3015. doi: 10.1681/ASN.2021070988
- Wang W, Zhu W, Hajagos J, Fochtmann L, Koraishy FM. Classification and regression trees analysis identifies patients at high risk for kidney function decline following hospitalization. PLoS One. 2025;20(1):e0317558. doi: 10.1371/journal.pone.0317558
- Poonia RC, Gupta MK, Abunadi I, et al. Intelligent diagnostic prediction and classification models for detection of kidney disease. Healthcare (Basel). 2022;10(2):371. doi: 10.3390/healthcare10020371
- Sabanayagam C, Banu R, Lim C, et al. Artificial intelligence in chronic kidney disease management: a scoping review. Theranostics. 2025;15(10):4566-4578. doi: 10.7150/thno.108552
- Nramban Kannan SK, Aseervatham J, Moholkar K, et al. A model for predicting chronic kidney diseases based on medical data using reinforcement learning. SN Comput Sci. 2024;5(4):353. doi: 10.1007/s42979-024-02665-z
- Hegde GM, Shenoy PD, Venugopal KR, Canchi A. A deep learning framework for chronic kidney disease stage classification. Healthc Anal. 2025;7:100398. doi: 10.1016/j.health.2025.100398
- Saif D, Sarhan AM, Elshennawy NM. Early prediction of chronic kidney disease based on ensemble of deep learning models and optimizers. J Electr Syst Inf Technol. 2024;11(1):17. doi: 10.1186/s43067-024-00142-4
- Bayor AA, Li J, Yang IA, Varnfield M. Designing clinical decision support systems (CDSS): a user-centered lens of the design characteristics, challenges, and implications: systematic review. J Med Internet Res. 2025;27:e63733. doi: 10.2196/63733
- Caruana R. Multitask learning. Mach Learn. 1997;28(1):41-75. doi: 10.1023/A:1007379606734
- Zhang Y, Yang Q. A survey on multi-task learning. IEEE Trans Knowl Data Eng. 2022;34(12):5586-5609. doi: 10.1109/TKDE.2021.3070203
- Chen J, Li Q, Liu F, Wen Y. M3T-LM: a multi-modal multi-task learning model for jointly predicting patient length of stay and mortality. Comput Biol Med. 2024;183:109237. doi: 10.1016/j.compbiomed.2024.109237
- Kim H, Lee SW, Kim SJ, et al. Advancing continuous prediction for acute kidney injury via multi-task learning: towards better clinical applicability. IEEE J Biomed Health Inform. 2025;29(8):5949-5962. doi: 10.1109/JBHI.2025.3559677
- Kendall A, Gal Y, Cipolla R. Multi-task learning using uncertainty to weigh losses for scene geometry and semantics. In: Proceedings of the 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). IEEE; 2018:7482-7491. doi: 10.1109/CVPR.2018.00781
- Lundberg SM, Lee SI. A unified approach to interpreting model predictions. In: Advances in Neural Information Processing Systems 30 (NeurIPS 2017). 2017:4765-4774.
- van Buuren S, Groothuis-Oudshoorn K. mice: multivariate imputation by chained equations in R. J Stat Softw. 2011;45(3):1-67. doi: 10.18637/jss.v045.i03
- Troyanskaya O, Cantor M, Sherlock G, et al. Missing value estimation methods for DNA microarrays. Bioinformatics. 2001;17(6):520-525. doi: 10.1093/bioinformatics/17.6.520
- Kapoor S, Narayanan A. Leakage and the reproducibility crisis in machine-learning-based science. Patterns. 2023;4(9):100804. doi: 10.1016/j.patter.2023.100804
- Singh V, Asari VK, Rajasekaran R. A deep neural network for early detection and prediction of chronic kidney disease. Diagnostics. 2022;12(1):116. doi: 10.3390/diagnostics12010116
- Kumar A, Sinha N, Bhardwaj A, Goel S. Clinical risk assessment of chronic kidney disease patients using genetic programming. Comput Methods Biomech Biomed Engin. 2022;25(8):887-895. doi: 10.1080/10255842.2021.1985476
- Swain D, Mehta U, Bhatt A, et al. A robust chronic kidney disease classifier using machine learning. Electronics. 2023;12(1):212. doi: 10.3390/electronics12010212
- Venkatrao K, Kareemulla S. HDLNET: a hybrid deep learning network model with intelligent IOT for detection and classification of chronic kidney disease. IEEE Access. 2023;11:99638-99652. doi: 10.1109/ACCESS.2023.3312183
- Arumugham V, Sankaralingam BP, Jayachandran UM, Krishna KVSSR, Sundarraj S, Mohammed M. An explainable deep learning model for prediction of early-stage chronic kidney disease. Comput Intell. 2023;39(6):1022-1038. doi: 10.1111/coin.12587
- Halder RK, Uddin MN, Uddin MA, et al. ML-CKDP: machine learning-based chronic kidney disease prediction with smart web application. J Pathol Inform. 2024;15:100371. doi: 10.1016/j.jpi.2024.100371
- Elshewey AM, Selem E, Abed AH. Improved CKD classification based on explainable artificial intelligence with extra trees and BBFS. Sci Rep. 2025;15(1):17861. doi: 10.1038/s41598-025-02355-7
- Ramu K, Patthi S, Prajapati YN, et al. Hybrid CNN-SVM model for enhanced early detection of chronic kidney disease. Biomed Signal Process Control. 2025;100:107084. doi: 10.1016/j.bspc.2024.107084
- Chowdhury MNH, Bin Ibne Reaz M, Ali SHM, et al. Deep learning for early detection of chronic kidney disease stages in diabetes patients: a TabNet approach. Artif Intell Med. 2025;166:103153. doi: 10.1016/j.artmed.2025.103153
- Arik SÖ, Pfister T. TabNet: attentive interpretable tabular learning. Proc AAAI Conf Artif Intell. 2021;35(8):6679-6687. doi: 10.1609/aaai.v35i8.16826
- Tangri N, Stevens LA, Griffith J, et al. A predictive model for progression of chronic kidney disease to kidney failure. JAMA. 2011;305(15):1553-1559. doi: 10.1001/jama.2011.451
- Little RJA. A test of missing completely at random for multivariate data with missing values. J Am Stat Assoc. 1988;83(404):1198-1202. doi: 10.1080/01621459.1988.10478722
- Kingma DP, Ba J. Adam: a method for stochastic optimization. arXiv. Preprint posted online 2014. doi: 10.48550/arXiv.1412.6980
- Vickers AJ, Elkin EB. Decision curve analysis: a novel method for evaluating prediction models. Med Decis Making. 2006;26(6):565-574. doi: 10.1177/0272989X06295361
- Van Calster B, McLernon DJ, van Smeden M, Wynants L, Steyerberg EW. Calibration: the Achilles heel of predictive analytics. BMC Med. 2019;17(1):230. doi: 10.1186/s12916-019-1466-7
- Lundberg SM, Erion G, Chen H, et al. From local explanations to global understanding with explainable AI for trees. Nat Mach Intell. 2020;2(1):56-67. doi: 10.1038/s42256-019-0138-9
- Heerspink HJL, Greene T, Tighiouart H, et al. Change in albuminuria as a surrogate endpoint for progression of kidney disease: a meta-analysis of treatment effects in randomised clinical trials. Lancet Diabetes Endocrinol. 2019;7(2):128-139. doi: 10.1016/S2213-8587(18)30314-0
- Babitt JL, Berns JS, Bozkurt B, et al. Executive summary of the KDIGO 2026 clinical practice guideline for the management of anemia in chronic kidney disease (CKD). Kidney Int. 2026;109(1):44-56. doi: 10.1016/j.kint.2025.06.005
- Tuttle KR, Agarwal R, Alpers CE, et al. Molecular mechanisms and therapeutic targets for diabetic kidney disease. Kidney Int. 2022;102(2):248-260. doi: 10.1016/j.kint.2022.05.012
- Ramspek CL, Jager KJ, Dekker FW, Zoccali C, van Diepen M. External validation of prognostic models: what, why, how, when and where? Clin Kidney J. 2021;14(1):49-58. doi: 10.1093/ckj/sfaa188
