Using Machine Learning to Develop Diagnostic Models for Cognitive Diseases

Abstract

Vascular Cognitive Disease (VCD), including Subcortical Small Vessel Disease (SSVD), remains one of the most underdiagnosed cases of dementia. Given the clinical need for early, accurate, and explainable diagnosis, this study explores machine learning techniques trained on real-world clinical data leveraged from the Gothenburg Mild Cognitive Impairment (MCI) study to uncover key variables in VCD diagnosis. The methodology was structured into three key parts: data preprocessing, model training and evaluation, and model explainability. The original dataset was partitioned into three subsets to reflect distinct clinical settings: primary care, specialist care, and research data, each with varying levels of feature availability. Given the high rate of missing values, multiple imputation techniques were explored and assessed. The model training part involved evaluating the performance of various machine learning algorithms across specific diagnostic tasks and clinical settings. These machine learning algorithms included two gradient boosting tree algorithms XGBoost and LightGBM, Support Vector Machine (SVM), Random Forest (RF), and Naive Bayes. The model with the highest average F1 score (ranging from 0 to 1) across multiple iterations was selected for final deployment and further refined through additional training. An Explainable AI (XAI) approach, named SHapley Additive exPlanations (SHAP), was applied to the final model to ensure transparency and clinical relevance and identify the most influential features contributing to classification outcomes. The majority of the final models could classify diagnoses with high precision and recall, achieving high F1 scores. Some of the variables were previously known to be associated with the diseases. Furthermore, new variables not previously linked to the disease were identified, prompting further research. In conclusion, the machine learning pipeline built in this study has the potential to act as a classifier to distinguish VCD from AD and clinical prestages of cognitive impairment in a clinical setting. Furthermore, it can be utilized to identify key variables associated with distinguishing between these diseases.