BrainSegFounder: Towards 3D Foundation Models for Neuroimage Segmentation

TL;DR

This paper presents BrainFounder, a novel 3D foundation model for neuroimage segmentation that leverages large-scale multi-modal MRI data and a two-stage pretraining approach to improve accuracy in neuroimaging tasks.

Contribution

Introduces BrainFounder, a new two-stage vision transformer-based model trained on extensive MRI data for improved neuroimage segmentation.

Findings

Outperforms previous state-of-the-art solutions on BraTS and ATLAS datasets.

Demonstrates the benefits of large-scale unlabeled data and model complexity.

Shows significant accuracy improvements in complex neuroimaging tasks.

Abstract

The burgeoning field of brain health research increasingly leverages artificial intelligence (AI) to interpret and analyze neurological data. This study introduces a novel approach towards the creation of medical foundation models by integrating a large-scale multi-modal magnetic resonance imaging (MRI) dataset derived from 41,400 participants in its own. Our method involves a novel two-stage pretraining approach using vision transformers. The first stage is dedicated to encoding anatomical structures in generally healthy brains, identifying key features such as shapes and sizes of different brain regions. The second stage concentrates on spatial information, encompassing aspects like location and the relative positioning of brain structures. We rigorously evaluate our model, BrainFounder, using the Brain Tumor Segmentation (BraTS) challenge and Anatomical Tracings of Lesions After Stroke v2.0 (ATLAS v2.0) datasets. BrainFounder demonstrates a significant performance gain, surpassing the achievements of the previous winning solutions using fully supervised learning. Our findings underscore the impact of scaling up both the complexity of the model and the volume of unlabeled training data derived from generally healthy brains, which enhances the accuracy and predictive capabilities of the model in complex neuroimaging tasks with MRI. The implications of this research provide transformative insights and practical applications in healthcare and make substantial steps towards the creation of foundation models for Medical AI. Our pretrained models and training code can be found at https://github.com/lab-smile/GatorBrain.