Fusing multimodal neuroimaging data with a variational autoencoder

Let $\{M_{i}=\{x_{i,j},...,x_{i,N}\}\}_{i=1,...,n}$ be a set of modalities with $N$ samples per modality and $n$ modalities. Sample index $j$ corresponds to a subject and each subject will have a sample for each modality. Instead of learning each modality with a separate decoder, we enforce a shared subspace. Further, to make this approach scalable to a large number of modalities and because we already use a shared decoder, we also only use one encoder for all modalities. This forces the features that are learned for each modality to be similar and makes sure the model scales well in terms of memory usage. Further, given that neuroimaging datasets are considered small compared to more commonly used deep learning datasets, using multiple encoders may lead to overfitting. The encoder decoder couple is optimized with respect to the log-likelihood of the marginal distribution of $M$ and each volume is treated as an independent sample. The integral over the marginal distribution of $M$, $p_{\theta}(x)=\int p_{\theta}(z)p_{\theta}(x|z)\mathit{d}z$ is intractable. We therefore optimize the following evidence lower bound (ELBO) using a variational autoencoder (VAE) [9]:

The objective function is calculated as a sum over each data point and can be split up into three parts, the first part is an encoder $q_{\phi}(z|x_{i,j})$ parameterized as a convolutional neural network (CNN) with parameters $\phi$ that estimates the latent variable $z$. The second part is a decoder $p_{\theta}(x_{i,j}|z)$, also parameterized as a CNN with parameters $\theta$, that reconstructs the original sample $x_{i,j}$ from the estimated latent variable $z$. The last part of the loss is the KL-divergence between a prior of our choosing $p_{\theta}(z)$, which we choose to be a diagonal multivariate Gaussian centered at $0$ with a standard deviation of $1$, and the estimated latent variable $z$. The latent variable $z$ is sampled from a multivariate Gaussian as well, which in turn is parameterized by a mean $\mu$ and variance $\sigma$ that is estimated by the encoder.

To evaluate whether the fusion of modalities in the VAE’s latent space leads to robust biomarkers, we set up a classification task. The model is first trained using 10-fold cross-validation, where each fold of subjects is used as a test set once and the other 9 folds are used to train on. The validation set is randomly selected as a stratified 10% of the subjects in the training set. After training the VAE in an unsupervised manner, the weights in the VAE are frozen. The complete dataset is then embedded using the encoder $q_{\phi}(z|x_{i,j})$, where instead of sampling $z$ from its estimated multivariate Gaussian, we use the estimated mean $\mu$ as our latent variable $z$. This is to make sure there is no stochasticity in the inference process. The representations of the training and validation sets are stacked and used as input for a machine learning model, this model is then evaluated using the test set representations, as follows.

The estimated latent variable $z$ can be interpreted as a low-dimensional representation of a volume, with a dimensionality $l$. Given that each subject has $n$ different modalities, each subject will also have $n$ representations $z_{1},...,n$. These representations can be concatenated for a subject to create a feature vector with a size of $n\times l$. The subject-by-feature matrix can be used as input for a classifier. In this case, we train a support vector machine (SVM)

to predict whether subjects in a held-out test set are patients with schizophrenia. Given that each modality is represented using $l$ features, we can extract the feature importance for all $nl$ features and then sum the features for each modality, to get feature importance for each of the $n$ modalities. The feature importance information helps us understand how brain changes related to schizophrenia are jointly represented in multiple modalities.

The SVM classifier is evaluated by calculating the area under the curve (AUC) of their receiver operator characteristic (ROC). Given that the VAE is trained in an unsupervised manner and that parts of the training process are stochastic, we evaluate the model for $5$ different seeds to make sure the method is robust to its seed, these experiments are performed with a latent dimensionality of $128$. To evaluate the effect of the number of latent dimensions on the classification performance, we set the seed to be $42$ and trained the model with four different latent dimensionalities $128$, $256$, $512$, $1024$. We select the weights obtained during the first training fold for the model that performs the best. The best performance is determined by averaging the ROC-AUC over the 10 folds for each model. These encoder and decoder weights are used to create the figures and to determine the importance of each modality for the classification of schizophrenia.

The datasets used in this study are FBIRN, B-SNIP, and COBRE, each dataset was processed using NeuroMark [11] to obtain 53 independent component networks (ICNs). These 53 ICNs, together with a structural MRI scan for each subject are considered to be separate modalities, so $n=54$. The sMRI data is preprocessed using SPM 12 in a Matlab 2016 environment. The data is then segmented into gray matter volume (GMV) with the help of a modulated algorithm, the GMV is then smoothed with a $6$mm FWHM Gaussian kernel.
Each ICN is a volume with $53\text{-by-}63\text{-by-}52$ voxels, the sMRI volumes are resized to that same size using Scipy [14]. The values in each volume are then rescaled to [-1, 1] by dividing the values in a volume by their maximum, which is also sometimes referred to as maximum absolute scaling. The dataloader and transformations were implemented with the help of TorchIO [15].

The batches are constructed by loading the $53$ ICN volumes and an sMRI volume for a subject and concatenating them into a batch. Each volume is considered to be sampled independently during training, the reason the volumes are loaded per subject however is to minimize disk accesses. The ICNs for one subject are all saved in one file, so loading all of them into a batch leads to a smaller number of disk accesses and reduces training time. Further, because each modality is equally present in a mini-batch for which an optimizer step is performed, the loss calculated over each batch is balanced for the modalities.

The code for the model, inference, and training was implemented using PyTorch [16], Catalyst [17], and NumPy [18]. The VAE uses a convolutional encoder and decoder pair, each of the layers uses a $3$-voxel kernel, a stride of $2$ and $1$-voxel padding. The channel sizes in the encoder are $1\to 64,64\to 128,128\to 256,256\to 512$ and $512\to 256,256\to 128,128\to 64,64\to 32,32\to 16,16\to 1$ in the decoder, the last layer in the decoder uses stride $1$ and no bias parameters. Each convolutional layer uses a ReLU [19] as its activation function, except for the last layer in the decoder, which uses a hyperbolic tangent function to map the output between [-1, 1] to match the input range. The last convolutional layer in the encoder produces an output with shape: $4\text{-by-}4\text{-by-}4$ and $256$ channels, this output is flattened and mapped to the mean $\mu$ and variance $\sigma$, which are used to construct a diagonal multivariate Gaussian from which $z$ is sampled. To make sure the VAE is fully differentiable, we use the reparameterization trick to train it [9]. The classification evaluations in the latent space are implemented using RAPIDS AI [20] to make sure highly parallelizable computations are performed on the GPU and to minimize costly CPU$\to$GPU and GPU$\to$CPU data transfers. The experiments were performed on an NVIDIA DGX-1 V100.

Most of the modalities that are used in this paper are intrinsic networks, which are obtained through independent component analysis (ICA). The independence in the spatial volumes for those components leads to a latent space that clusters modalities, which is shown in Figure 1. The plot depicts a T-SNE projection of a $512$-dimensional space that preserves the latent structure in a 2D image. Interestingly, the ICNs that belong to the same domain are generally clustered together, such as ICNs in the cerebellum. It is also clear from Figure 1 that the sMRI cluster is located relatively farther away from the other modalities in the latent space. The ICNs are all localized spatial ICA maps, whereas the sMRI volume represents all of the structures in the brain. There is more inter-subject variance to be modeled for the sMRI volumes than for the spatially localized ICNs which likely contributes to the sMRI cluster being further away from the latent space ICN clusters.

The average ROC-AUC for the five models trained with a latent dimensionality of $128$ and multiple different seeds is $0.8374$, with a standard deviation of $0.0026$. This shows that the model robustly learns a latent space, where patients with schizophrenia and controls are linearly separable.

The experiment involving an increasing number of latent dimensions shows that the ROC-AUC increases with the number of latent dimensions up to a certain point. The ROC-AUC is $0.8354$ for $128$ dimensions, $0.8569$ for $256$ dimensions, $0.8610$ for $512$ dimensions and $0.8540$ for $1024$ dimensions. These results show that the model is not learning to represent each modality using a single latent dimension, but rather that the variance that is modeled across a latent dimension is shared among multiple modalities or that multiple latent dimensions are used to model a single modality.

The best model, $seed=42$ and $512$ latent dimensions were used to calculate the feature importance for each modality. The 10 modalities with the highest feature importance are shown in Figure 2, where the rightmost modality is the most important and the leftmost modality is the 10th most important. sMRI shows the lowest performance in terms of feature importance. This is likely due to the trade-off in the loss function for the VAE. The KL-divergence pulls the latent variables closer to a zero-mean unit-norm Gaussian, while the reconstruction loss tries to make sure every modality is reconstructed correctly. The number of different modalities and the KL-divergence likely limit variance that can be modeled to represent the sMRI volumes. The variations that are modeled for sMRI do not help linearly separate patients from controls.

The group differences that the VAE has learned can be interpreted by visualizing the group centers in the latent space. Group centers can be calculated for each modality by averaging the locations of subjects within that group. The latent center for SZ patients can then be decoded and subtracted from the decoded latent center for HC to show the group differences. The differences of the top five most important features from Figure 2 are calculated, then thresholded at the 99th quantile highest values for each modality, and then summed to create Figure 3. The figure compares the learned differences with the voxelwise differences of the spatial ICA components that correspond to the five most important modalities. The results are remarkably similar, which proofs that VAEs yield interpretable results in their latent space. As future work fuses modalities more, the results will likely differ more from the voxelwise differences.