InfoVAE-Med3D: Interpretable 3D MRI

• The paper demonstrates that InfoVAE-Med3D maximizes mutual information between 3D brain MRIs and latent codes, preserving key anatomical features for high-fidelity reconstruction.
• It introduces a novel loss function that balances reconstruction fidelity with latent space regularity through explicit mutual information and KL divergence terms.
• Results show enhanced predictive accuracy for brain age and cognitive scores and improved interpretability over traditional VAE methods.

InfoVAE-Med3D is a latent representation learning framework for 3D medical imaging that extends InfoVAE to maximize mutual information between volumetric images and latent codes, thereby enabling compact, interpretable, and clinically meaningful embeddings. Designed originally for 3D brain MRI, InfoVAE-Med3D unites high-fidelity image reconstruction with downstream predictive tasks such as brain-age and cognitive score regression, outperforming conventional VAEs and offering insights into biomarker discovery and neurological disease analysis.

1. Theoretical Foundations and Objective Formulation

InfoVAE-Med3D is based on the InfoVAE paradigm, which generalizes variational autoencoder (VAE) objectives by introducing explicit mutual information terms into the loss function. Formally, for a 3D brain MRI input $X$ and latent variable $Z$, the model aims to maximize the mutual information $MI_q(X;Z)$ while maintaining reconstruction fidelity and latent space regularity:

$$\mathcal{L}_{\text{InfoVAE-Med3D}} = \mathbb{E}_{p(X)}[\mathcal{L}_{\text{rec}}(X)] - \alpha \cdot MI_q(X;Z) - \beta \cdot D_{KL}(q_\phi(Z) \parallel p(Z))$$

where:

• $\mathcal{L}_{\text{rec}}(X)$ is the per-sample reconstruction loss, typically measured via $\ell_1$, $\ell_2$, or perceptual similarity metrics.
• $$MI_q(X;Z) = \mathbb{E}_q [\log \big( q_\phi(X, Z) / (q_\phi(X) \cdot q_\phi(Z)) \big)]$$ quantifies the information retained about $X$ in $Z$.
• $D_{KL}(q_\phi(Z) \| p(Z))$ ensures the aggregate posterior's alignment with the prior.
• $\alpha$ and $\beta$ are tuning hyperparameters regulating mutual information and regularization strengths.

This formulation counteracts posterior collapse and drives the latent space to encode rich, structured, and clinically relevant anatomical content.

2. Data Sources and Evaluation Protocols

InfoVAE-Med3D has been evaluated on two major cohorts:

• Healthy Control Dataset ("BrainAge"): 6,527 subjects, age range 18–97 (mean 43.67), with paired 3D MRI and gender. Train/val/test split: 5,221/653/653.
• Multiple Sclerosis (MS) Dataset (Charles University, Prague): 904 patients, paired MRI, chronological age (mean ~42.2), Symbol Digit Modalities Test (SDMT, mean ~58.9). Typical split: 733/95/88.

Evaluation encompasses:

• Reconstruction: Using both quantitative (SSIM, PSNR) and qualitative assessment across views (axial, sagittal, coronal).
• Downstream Regression: Brain age and SDMT scores are predicted via support vector regression on latent codes, measured by mean absolute error (MAE), $R^2$, and root mean squared error (RMSE).
• Latent Space Visualization: Clustering and gradient analysis with PCA and supervised partial least squares (PLS) dimension reduction reveal interpretability.

3. Experimental Results

Across tasks, InfoVAE-Med3D demonstrates improved fidelity and predictive capacity over VAE and $\beta$-VAE baselines:

• Reconstruction: For example, SSIM = 0.750, PSNR = 24.91 (BrainAge, $\alpha = 0$, $\beta = 1$), outperforming both AE and classical VAE.
• Brain Age Regression: Reduced MAE and RMSE, elevated $R^2$, surpassing competitors.
• SDMT Prediction: Enhanced accuracy, despite limited direct label association within latent space.

Qualitative results confirm that anatomical structures—e.g., cortex, hemisphere separation—are better preserved than in baseline approaches. InfoVAE-Med3D achieves sharper, clinically faithful reconstructions and latent representations that reflect underlying phenotypes.

4. Interpretability via Latent Space Analysis

A core goal is interpretable embeddings:

• Incorporating $MI_q(X;Z)$ guarantees retention of essential image features.
• Latent vectors (e.g., 512D) are projected to 2D:
  • PCA visualizes population variability.
  • PLS enhances attribute-driven clusters (age, gender, cognitive scores).
• Post-training, clusters manifest clinically coherent separations: e.g., gender-based grouping, age-driven gradients, and finer subtleties for cognitive score (SDMT).

Such visualization allows direct inspection of disease-relevant features and supports trust in model outputs.

5. Methodological Considerations and Trade-offs

The InfoVAE-Med3D approach manages a computational and statistical balance:

• Statistical: Tuning $\alpha$ and $\beta$ regulates the density and regularity of the latent space, balancing reconstruction quality against interpretability.
• Computational: The explicit mutual information term introduces additional estimation cost versus standard ELBO-based VAEs. However, the approach maintains tractability for high-dimensional volumetric data.
• Model Structure: The encoder and decoder networks are adapted for 3D volumes (e.g., 3D convolutions), and the loss terms are coordinated to avoid latent code saturation or collapse.

A plausible implication is that parameter choices directly affect downstream biomarker discovery—higher $\alpha$ preserves more anatomical detail, whereas stronger $\beta$ aligns the latent manifold for generative modeling and cohort-level analyses.

6. Clinical and Research Applications

InfoVAE-Med3D enables several key applications in neuroimaging:

• MRI-Based Biomarkers: Latent variables learned via InfoVAE-Med3D can serve as interpretable, quantitative indices for early neurodegeneration or cognitive impairment.
• Disease Analysis: Regression against brain age and SDMT supports the noninvasive evaluation of cognitive decline, with direct mapping from high-dimensional MRI features to clinical outcomes.
• Research and Personalized Medicine: The approach facilitates advanced cohort stratification, exploration of structural–functional relationships, and personalized monitoring.
• Transparency in Deep Learning: Structured, interpretable latent spaces foster clinician trust and support integration into diagnostic workflows.

7. Relationship to Broader InfoVAE and Medical 3D Frameworks

InfoVAE-Med3D is a direct instantiation of information-theoretic latent variable modeling applied to medical 3D imaging, consistent with the Lagrangian dual view articulated in the InfoVAE family (Zhao et al., 2018). Its explicit mutual information maximization extends theoretical results connecting interpretability, reconstruction, and regularization. The methodological innovations in InfoVAE-Med3D complement, and in some settings surpass, alternative VAE-based approaches in both predictive accuracy and transparency.

This suggests that InfoVAE-Med3D constitutes a practical, scalable tool for interpretable latent representation learning in neuroimaging and related volumetric biomedical domains.