BrainSegFounder: Towards 3D foundation models for neuroimage segmentation.

in Medical image analysis by Joseph Cox, Peng Liu, Skylar E Stolte, Yunchao Yang, Kang Liu, Kyle B See, Huiwen Ju, Ruogu Fang

TLDR

  • The study develops a new way to analyze brain images using artificial intelligence. The approach involves training a model using unlabeled images of healthy brains and then fine-tuning the model to identify specific features in diseased brains. The study shows that this approach outperforms previous methods and can be used to diagnose and plan treatments for various brain diseases.

Abstract

The burgeoning field of brain health research increasingly leverages artificial intelligence (AI) to analyze and interpret neuroimaging data. Medical foundation models have shown promise of superior performance with better sample efficiency. This work introduces a novel approach towards creating 3-dimensional (3D) medical foundation models for multimodal neuroimage segmentation through self-supervised training. Our approach involves a novel two-stage pretraining approach using vision transformers. The first stage encodes anatomical structures in generally healthy brains from the large-scale unlabeled neuroimage dataset of multimodal brain magnetic resonance imaging (MRI) images from 41,400 participants. This stage of pertaining focuses on identifying key features such as shapes and sizes of different brain structures. The second pretraining stage identifies disease-specific attributes, such as geometric shapes of tumors and lesions and spatial placements within the brain. This dual-phase methodology significantly reduces the extensive data requirements usually necessary for AI model training in neuroimage segmentation with the flexibility to adapt to various imaging modalities. We rigorously evaluate our model, BrainSegFounder, using the Brain Tumor Segmentation (BraTS) challenge and Anatomical Tracings of Lesions After Stroke v2.0 (ATLAS v2.0) datasets. BrainSegFounder 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 model complexity and the volume of unlabeled training data derived from generally healthy brains. Both of these factors enhance the accuracy and predictive capabilities of the model in neuroimage segmentation tasks. Our pretrained models and code are at https://github.com/lab-smile/BrainSegFounder.

Overview

  • The study aims to develop a novel approach for creating 3D medical foundation models for multimodal neuroimage segmentation through self-supervised training using vision transformers. The approach involves a two-stage pretraining process that identifies key features in healthy brains and disease-specific attributes in diseased brains. The study evaluates the performance of the model, BrainSegFounder, using the BraTS challenge and ATLAS v2.0 datasets and demonstrates a significant performance gain compared to previous winning solutions using fully supervised learning. The pretrained models and code are available at <https://github.com/lab-smile/BrainSegFounder>.

Comparative Analysis & Findings

  • The study compares the performance of BrainSegFounder, a self-supervised learning model, with previous winning solutions using fully supervised learning on the BraTS challenge and ATLAS v2.0 datasets. The results show that BrainSegFounder outperforms the previous solutions, achieving a significant performance gain. The study also identifies the importance of scaling up both the model complexity and the volume of unlabeled training data derived from generally healthy brains in enhancing the accuracy and predictive capabilities of the model in neuroimage segmentation tasks.

Implications and Future Directions

  • The study's findings highlight the potential of self-supervised learning and vision transformers in improving the accuracy and efficiency of neuroimage segmentation tasks. The approach can be adapted to various imaging modalities and can be used for disease diagnosis and treatment planning. Future research directions include exploring the use of BrainSegFounder for other neuroimaging tasks, such as connectomics and functional connectivity analysis, and investigating the impact of different pretraining strategies on model performance.
Read Full Article