Variational voxelwise rs-fMRI representation learning: Evaluation of sex, age, and neuropsychiatric signatures

This work uses a variational autoencoder (VAE) to perform non-linear representation learning from voxelwise rs-fMRI data. The VAE learns a non-linear dimensionality reduction of the data in the form of a latent vector. These latent vectors retain meaningful information related to a subject’s demographics and clinical diagnosis. The retention of meaningful information in the latent vectors is evaluated using age regression and sex classification tasks on the UK Biobank dataset. The results on these tasks are highly encouraging and a linear regressor trained on the latent vectors to predict age performs almost on par with a supervised neural network. Further, the same latent vectors can almost perfectly linearly separate sex. The model that is pre-trained on UK Biobank is also fine-tuned on a smaller neuropsychiatric dataset for a varying number of epochs. The latent vectors it generates for this dataset are then evaluated by performing a schizophrenia diagnosis classification task. We find that pre-training the model on UK Biobank significantly improves the quality of the latent vectors and that the vectors themselves are fairly discriminative. To understand the structure of the latent vectors with respect to demographic variables or neuropsychiatric disorders we train a variety of supervised models on the latent vectors. The results presented in this work open up more in-depth research into the factors of variation that the VAE models and how they can be improved for voxelwise rs-fMRI data.