Latent Representation Learning from 3D Brain MRI for Interpretable Prediction in Multiple Sclerosis

TL;DR

The paper tackles the challenge of linking 3D brain MRI to cognitive decline biomarkers by learning structured latent embeddings that are both predictive and interpretable. It extends the InfoVAE framework to maximize $I(X;Z)$, producing compact 512-dimensional latent vectors from 3D MRI and enabling brain-age and SDMT regression while preserving clinically meaningful content. Across BrainAge and Prague MS cohorts, InfoVAE-Med3D outperforms standard VAE variants in reconstruction and downstream tasks, and its latent space reveals interpretable patterns such as gender clustering and smooth age trajectories. This approach provides a practical path toward MRI-based biomarkers with transparent analysis of cognitive deterioration, with potential to advance biomarker discovery in MS and related neurodegenerative conditions.

Abstract

We present InfoVAE-Med3D, a latent-representation learning approach for 3D brain MRI that targets interpretable biomarkers of cognitive decline. Standard statistical models and shallow machine learning often lack power, while most deep learning methods behave as black boxes. Our method extends InfoVAE to explicitly maximize mutual information between images and latent variables, producing compact, structured embeddings that retain clinically meaningful content. We evaluate on two cohorts: a large healthy-control dataset (n=6527) with chronological age, and a clinical multiple sclerosis dataset from Charles University in Prague (n=904) with age and Symbol Digit Modalities Test (SDMT) scores. The learned latents support accurate brain-age and SDMT regression, preserve key medical attributes, and form intuitive clusters that aid interpretation. Across reconstruction and downstream prediction tasks, InfoVAE-Med3D consistently outperforms other VAE variants, indicating stronger information capture in the embedding space. By uniting predictive performance with interpretability, InfoVAE-Med3D offers a practical path toward MRI-based biomarkers and more transparent analysis of cognitive deterioration in neurological disease.

Figures (4)

• Figure 1: Qualitative comparison of InfoVAE-Med3D with three VAE variants across coronal, sagittal, and axial views. Two examples are shown: a BrainAge sample (left, blue background) and a Prague sample (right, yellow background).
• Figure 2: Two-dimensional visualization of latent representations colored by SDMT scores. The PCA projection (left) shows partial separation, while the PLS regression projection (right) reveals an improved SDMT gradient but still not clearly defined.
• Figure 3: Two-dimensional visualization of latent representations colored by gender labels, obtained using the first two components of each method. Each subfigure presents PCA on the left and PLS regression on the right.
• Figure 4: Two-dimensional visualization of latent representations colored by age values, obtained using the first two components of each method. Each subfigure presents PCA on the left and PLS regression on the right.