# Comparative machine learning to predict acute kidney injury in traumatic brain injury: a MIMIC-IV cohort with SHAP interpretation

## TL;DR

This study uses machine learning to predict acute kidney injury in traumatic brain injury patients, finding that XGBoost performs best with strong accuracy and clinical relevance.

## Contribution

The study introduces a robust XGBoost model with SHAP interpretability for early AKI prediction in TBI patients using MIMIC-IV data.

## Key findings

- XGBoost outperformed other models with an AUC of 0.775 and high sensitivity (88.3%).
- SHAP analysis highlighted urine output and ventilation as key predictors of AKI in TBI patients.
- Ensemble models like XGBoost and RF showed better performance than traditional logistic regression.

## Abstract

AKI is a frequent and severe complication among TBI patients. Accurate early prediction is critical but remains challenging in ICU practice.

We retrospectively analyzed the MIMIC-IV database. After screening 85,242 first ICU admissions and applying exclusions, 2,986 TBI patients were included. AKI was defined by KDIGO criteria. Demographic, physiological, laboratory, and intervention variables were extracted, preprocessed, and imputed. Predictors were selected using LASSO, Boruta, and logistic regression with bootstrap validation. Seven ML models (LR, DT, RF, XGBoost, LightGBM, SVM, ANN) were trained on 70% of the cohort and validated on 30%, with hyperparameters optimized by grid search and 5-fold CV. Performance was assessed by AUC, calibration, DCA, accuracy, sensitivity, specificity, PPV, NPV, and F1-score. SHAP was applied to the best-performing model (XGBoost) for global and individual interpretability.

Of the 2,986 TBI patients, 2,045 (68.5%) developed AKI. AKI patients were older, heavier, and had higher glucose, sodium, SBP, and temperature, with lower urine output and more frequent ventilation. Feature selection consistently retained urine output, ventilation, weight, age, glucose, sodium, SBP, and temperature as core predictors. In validation, XGBoost showed the best performance (AUC 0.775, 95% CI 0.747–0.802; accuracy 74.4%; sensitivity 88.3%; F1-score 0.83), followed by RF (AUC 0.768; sensitivity 85.9%; F1-score 0.82). LR had moderate discrimination (AUC 0.763) but poor specificity (36.5%), while LightGBM achieved the highest specificity (50.4%) but lower AUC (0.741). DT performed worst (AUC 0.728; accuracy 69.3%). Calibration and DCA supported XGBoost as having the greatest clinical benefit. SHAP analysis of XGBoost identified urine output and ventilation as dominant predictors and provided patient-level explanations consistent with observed clinical patterns.

Ensemble ML models, particularly XGBoost, demonstrated robust predictive power, outperforming LR and DT. The XGBoost model combined high discrimination, calibration, and interpretability, offering a clinically applicable tool for early AKI risk stratification in TBI.