Simultaneous Prediction of Multiple Functional Outcomes in Post-Stroke Rehabilitation Using Machine Learning

Abstract

Introduction

Stroke is a leading cause of long-term disability worldwide. Despite the central role of rehabilitation to functional recovery, accurately predicting outcomes remains difficult due to the marked variability in post-stroke recovery trajectories. Machine learning may help modelling complex, non-linear relationships between clinical variables and outcomes. However, most existing approaches focus on single outcomes, neglecting the multidimensional nature of post-stroke recovery. Multi-output models may represent a more clinically coherent and statistically efficient alternative.

Aims
This study aims to compare single-output and multi-output machine learning approaches for predicting rehabilitation outcomes at discharge in post-stroke patients undergoing intensive rehabilitation. Specifically, we trained and compared four algorithms in both their single-output and multi-output formulations to predict three functional outcomes and evaluated whether joint modelling improves prediction accuracy over independent single-output machine learning models.

Methods

Data were obtained from the STRATEGY registry, a multicentre Italian prospective cohort of 743 post-stroke patients admitted for intensive inpatient rehabilitation at IRCCS Fondazione Don Carlo Gnocchi centres from 2022 onward. The registry systematically collects standardised demographic, clinical, neurological, cognitive, and functional data from adults admitted within 30 days of ischemic or haemorrhagic stroke, assessed at admission and discharge according to the PMIC2020 framework and additional measures of clinical complexity and comorbidity. Outcomes of interest were Modified Barthel Index (mBI), Trunk Control Test (TCT), and modified Rankin Scale (mRS) scores at discharge. After data cleaning, 702 patients were included. Model development and evaluation were conducted using a nested validation scheme. In the outer loop, the dataset was split into a development set (n = 561, 80%) and a held-out test set (n = 141, 20%). Within the development set, an inner 5-fold cross-validation (CV) procedure was used for hyperparameter tuning and model selection. Missing data were imputed using k-Nearest Neighbours, with the imputation model fitted exclusively on training folds to prevent information leakage. Each training fold and the final development set were augmented via synthetic data generation using the R package synthpop.

Four machine learning algorithms were evaluated in both single-output and multi-output formulations: Random Forest, Elastic Net, CatBoost, and XGBoost. For each algorithm, hyperparameters were optimised through grid search within the inner 5-fold CV. Model selection was based on a composite metric defined as the average Normalised Root Mean Square Error (NRMSE) across the three outcomes. The best-performing single-output and multi-output models were retrained on the full development set and evaluated on the held-out test set.

Results

The best-performing multi-output model was the multi-output Random Forest, whereas Elastic Net achieved the best performance among single-output models, with comparable performance during cross-validation. On the held-out test set, the multi-output approach achieved better performance on all three outcomes, with the largest improvements observed for mBI and TCT.

Conclusions

Multi-output modelling, by leveraging the shared structure among rehabilitation outcomes, improved prediction accuracy across all three targets compared to independent single-output models. These findings suggest that multi-output machine learning approaches may provide a more effective framework for prognostic modelling in post-stroke rehabilitation.