(G-242) Large Brain Models: Large-scale 3D Pretrained Deep Learning Models for Neuroimages

Introduction:: Foundation models (FMs) are powerful artificial intelligence models trained on massive unlabeled datasets to generate versatile models that can be adapted to various downstream tasks such as language translation and medical image analysis. FMs have enabled the development of groundbreaking products such as ChatGPU and DALL-E, which leverage large-scale natural language and image pre-trained models. The healthcare industry has long struggled with limited labeled data and expensive labor costs for data collection and annotation, but the development of medical domain specific FMs has the potential to revolutionize the field. However, as many medical images are 3-dimensional or even 4-dimensional, they may not benefit from the current FMs that are pre-trained on 2-dimensional data. To advance medical artificial intelligence (MAI) in brain health, we propose the construction of a large-scale 3D pre-trained deep learning model for neuroimages.

Materials and Methods:: Utilizing 12,000 T1 and T2-FLAIR MRI volumes from the UK Biobank we created a large-scale neuroimage-pretrained vision transformer by implementing a modified Sliding window UNET Transformer (Swin UNETR) in 3D space. This model can be transferred to much smaller disease-specific datasets such as tumor detection, early detection of Alzheimer’s Disease (AD), and prediction of cognitive decline. To validate the effectiveness of the model, we fine-tuned the generated model on both the Brain Tumor Segmentation (BraTS) challenge and classification of images in the Alzheimer’s Disease, Alzheimer's Disease Neuroimaging Initiative (ADNI) database.


Results, Conclusions, and Discussions:: The final pretrained model was generated in 5 hours on 128 NVIDIA A100 GPUs. Our model outperforms the top performers in the BraTS challenge even when restricted to only two input modalities compared to the four utilized in the winning submissions, achieving a mean Dice score of 0.920 across 5 folds compared to the challenge winner’s nnU-Net Dice score of 0.908. The model trained from only T1 volumes still performs better than nn-UNet with a mean Dice score of 0.909. On the ADNI dataset, we obtained an accuracy of 97.6% in distinguishing AD from normal cognitive function, which compares favorably to UNETR without pretrained weights and DenseNet.

We present a neuroimage-based pretrained network that solves images in their native 3D space rather than requiring reformulation to 2D or 2.5D. Our results suggest that our large-scale pretrained neuroimage model fully leverage the 3D information interleaved in multiple MRI imaging modalities. We demonstrate that our model can be readily adapted for multiple downstream tasks.

Our top-scoring results on the BraTS challenge and state-of-the-art performance on the ADNI dataset validate the effectiveness of our pretraining framework and model architecture. Fine-tuning our model achieves higher accuracy and more efficient convergence than training from scratch. This model would enable faster and more accurate medical image analysis, diagnosis, and treatment planning, among other applications, ultimately improving patient outcomes and the overall quality of healthcare. Future work can be directed at fine-tuning on further downstream tasks and extending the model to more modalities.

Acknowledgements (Optional): : Funded by the NSF IIS

References (Optional): :