Disentangling Hippocampal Shape Variations: A Study of Neurological Disorders Using Mesh Variational Autoencoder with Contrastive Learning

TL;DR

This research provides valuable insights into the relationship between neurological disorder and hippocampal shape changes in different age groups of MS populations using a Mesh VAE with Supervised Contrastive loss.

Abstract

This paper presents a comprehensive study focused on disentangling hippocampal shape variations from diffusion tensor imaging (DTI) datasets within the context of neurological disorders. Leveraging a Mesh Variational Autoencoder (VAE) enhanced with Supervised Contrastive Learning, our approach aims to improve interpretability by disentangling two distinct latent variables corresponding to age and the presence of diseases. In our ablation study, we investigate a range of VAE architectures and contrastive loss functions, showcasing the enhanced disentanglement capabilities of our approach. This evaluation uses synthetic 3D torus mesh data and real 3D hippocampal mesh datasets derived from the DTI hippocampal dataset. Our supervised disentanglement model outperforms several state-of-the-art (SOTA) methods like attribute and guided VAEs in terms of disentanglement scores. Our model distinguishes between age groups and disease status in patients with Multiple Sclerosis (MS) using the hippocampus data. Our Mesh VAE with Supervised Contrastive Learning shows the volume changes of the hippocampus of MS populations at different ages, and the result is consistent with the current neuroimaging literature. This research provides valuable insights into the relationship between neurological disorder and hippocampal shape changes in different age groups of MS populations using a Mesh VAE with Supervised Contrastive loss. Our code is available at https://github.com/Jakaria08/Explaining_Shape_Variability

Figures (7)

• Figure 1: Overall architecture of our method. We have mesh VAE with an encoder $f_\phi(x)$ and decoder $f_\theta(z)$ where $x$ is the input 3D mesh and $z$ represents the latent space. $L_{vae}$ represents the VAE loss that combines reconstruction and KL divergence loss. Another two losses are classification loss $L_{contr}^{cls}$ and regression loss $L_{contr}^{reg}$, where a specific latent variable is disentangled for a specific feature (continuous or discrete). We use the first variable for contrastive classification loss ($z_{1}$ corresponds to binary labels, and the rest variables are uncorrelated to the labels). The second variable $z_{2}$ corresponds to regression loss, and the rest variables are uncorrelated to the continuous labels.
• Figure 2: On the left side, we show the combination of reconstructions and original torus meshes from the synthetic dataset using our proposed model. The dark blue indicates a very small deviation between the reconstruction and the original mesh. We show two variabilities in the matrix of images: bump height and scale. On the right side, we show the decoder’s output by varying the disentangled latent variable $z_{1}$ and $z_{2}$ in the x and y axis while holding the other latent variables constant at a mean value which is zero.
• Figure 3: On the left side of the figure, we show the combination of reconstructions and original hippocampus (left and right hippocampus) meshes from the dataset using our proposed model. The dark blue indicates a very small deviation between the reconstruction and the original mesh. On the right side, we show the original hippocampus data.
• Figure 4: Effect of Spiral Sequence Length on SAP Score for our Torus and Hippocampus Datasets.
• Figure 5: SAP scores (Classification and regression are separated) for different models using synthetic torus (left) and hippocampus (right) datasets.