Latent Representation Learning from 3D Brain MRI for Latent Analysis and Prediction in Multiple Sclerosis.

Authors

Abstract

Neurological diseases can lead to significant cognitive impairment. Linking magnetic resonance imaging (MRI)-derived measures of brain structure with clinical assessments of cognitive function remains challenging, as conventional statistical methods and traditional machine learning approaches may struggle to capture complex relationships between brain structure and cognitive outcomes, hindering the development of reliable imaging biomarkers. Although deep learning methods offer greater predictive capabilities, their black-box nature limits interpretability, posing a major obstacle to clinical translation. Latent representation learning via generative models can provide informative representations and support deeper analysis and cognitive biomarker prediction. While standard generative adversarial networks and diffusion models primarily focus on generating images from random noise, variational autoencoders offer compact, probabilistic latent representations that are well suited for downstream tasks. In this study, we propose InfoVAE-Med3D, an adaptation of the InfoVAE framework for three-dimensional brain ​MRI to learn informative latent representations. By explicitly relaxing the mutual information regularization penalty between inputs and latent representations, our method aims to capture richer information from input MRI data, facilitating the analysis of brain-related characteristics and downstream prediction. We evaluate InfoVAE-Med3D on two datasets: a healthy control dataset ([Formula: see text]) containing chronological age and sex information, and a Multiple Sclerosis (MS) dataset ([Formula: see text]) collected from the Multiple Sclerosis Center of the First Medical Faculty, Charles University, Prague, containing chronological age, sex, and Symbol Digit Modalities Test (SDMT) scores. The learned latent representations capture information related to sex and brain volumetric features, including grey and white matter volumes, as reflected in sex-related clustering and smooth transitions in brain volume measures in two-dimensional visualizations. Through comprehensive evaluations, we demonstrate that InfoVAE-Med3D learns more informative latent representations, especially when the mutual information regularization penalty is completely removed. In this setting, the model generally outperforms other variational autoencoder variants in both MRI reconstruction and prediction of brain age and SDMT scores using support vector regression, while remaining competitive with deep learning-based predictive models, particularly on the MS dataset. InfoVAE-Med3D provides informative latent representations from brain MRI data of both healthy controls and patients with MS. By improving predictive performance and providing interpretable insights into the learned latent representations, the proposed framework offers a promising approach for advancing biomarker discovery and the analysis of cognitive impairment in neurological diseases, particularly ​MS.

Tags

Topics