Efficiently Training Vision Transformers on Structural MRI Scans for Alzheimer's Disease Detection

TL;DR

This paper explores the application of vision transformers to 3D brain MRI scans for Alzheimer's disease detection, demonstrating high accuracy and the importance of training strategies in neuroimaging tasks.

Contribution

It introduces ViT variants for neuroimaging, evaluates training strategies like pre-training and data augmentation, and analyzes data amount effects on performance.

Findings

Achieved AUC of 0.987 for sex classification and 0.892 for AD classification.

Fine-tuning pre-trained ViT models improved performance by 5-10%.

Training strategies are crucial for effective ViT application in neuroimaging.

Abstract

Neuroimaging of large populations is valuable to identify factors that promote or resist brain disease, and to assist diagnosis, subtyping, and prognosis. Data-driven models such as convolutional neural networks (CNNs) have increasingly been applied to brain images to perform diagnostic and prognostic tasks by learning robust features. Vision transformers (ViT) - a new class of deep learning architectures - have emerged in recent years as an alternative to CNNs for several computer vision applications. Here we tested variants of the ViT architecture for a range of desired neuroimaging downstream tasks based on difficulty, in this case for sex and Alzheimer's disease (AD) classification based on 3D brain MRI. In our experiments, two vision transformer architecture variants achieved an AUC of 0.987 for sex and 0.892 for AD classification, respectively. We independently evaluated our models on data from two benchmark AD datasets. We achieved a performance boost of 5% and 9-10% upon fine-tuning vision transformer models pre-trained on synthetic (generated by a latent diffusion model) and real MRI scans, respectively. Our main contributions include testing the effects of different ViT training strategies including pre-training, data augmentation and learning rate warm-ups followed by annealing, as pertaining to the neuroimaging domain. These techniques are essential for training ViT-like models for neuroimaging applications where training data is usually limited. We also analyzed the effect of the amount of training data utilized on the test-time performance of the ViT via data-model scaling curves.