Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

BrainSegFounder: Towards 3D Foundation Models for Neuroimage Segmentation

Joseph Cox, Peng Liu, Skylar E. Stolte, Yunchao Yang, Kang Liu, Kyle B. See, Huiwen Ju, Ruogu Fang

TL;DR

The proposed model, BrainSegFounder, demonstrates a significant performance gain, surpassing the achievements of the previous winning solutions using fully supervised learning, and underscores the impact of scaling up both the model complexity and the volume of unlabeled training data derived from generally healthy brains.

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.

BrainSegFounder: Towards 3D Foundation Models for Neuroimage Segmentation

TL;DR

The proposed model, BrainSegFounder, demonstrates a significant performance gain, surpassing the achievements of the previous winning solutions using fully supervised learning, and underscores the impact of scaling up both the model complexity and the volume of unlabeled training data derived from generally healthy brains.

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.
Paper Structure (26 sections, 5 figures, 9 tables)

This paper contains 26 sections, 5 figures, 9 tables.

Table of Contents

  1. Introduction
  2. Data Acquisition and Preprocessing
  3. UK Biobank dataset
  4. BRaTS dataset
  5. ATLAS v2.0 Dataset
  6. Stage 1: Pretraining on the UKB
  7. Training on BraTS
  8. Stage 2: Pretraining on BraTS
  9. Stage 3: Fine-tuning on BraTS
  10. Few-shot Learning on BraTS
  11. Modality Restriction and Flexibility in Training and Inference
  12. Training on ATLAS v2.0
  13. Stage 2: Pretraining on ATLAS v2.0
  14. Stage 3: Fine-tuning on ATLAS v2.0
  15. Results
  16. ...and 11 more sections

Figures (5)

  • Figure 1: Overall Study Design. a) The two-stage pretraining process using Swin Transformer decoders and encoder. Initially, the model is pretrained on the UKB dataset (Stage 1), followed by the downstream task dataset (Stage 2). b) This is succeeded by fine-tuning on each downstream dataset, with transfer learning applied between each stage.
  • Figure 2: Visual representation of demographic data from subjects in the UK Biobank in the study.
  • Figure 3: CONSORT diagram of UKB data used in Stage 1 pretraining.
  • Figure 4: Training (left) and validation (right) loss of Stage 1-pretraining three different scale of BrainSegFounder models on UKB.
  • Figure 5: Dice coefficients for baseline (SwinUNETR) and our model across different levels of training data availability. All models were trained 5 times to account for variability in the input data randomly selected. Error bars represent $\pm$ one standard deviation.