Brain-SAM: A SAM-based Model Tailored for Brain MRI Lesion Segmentation

Technical Summary of Brain-SAM

Problem Statement
Magnetic resonance imaging (MRI) is fundamental to neuroimaging, yet accurate segmentation of brain structures and lesions remains challenging. While foundation models like the Segment Anything Model (SAM) and its successor SAM2 have demonstrated remarkable generalizability in natural images, their application to brain MRI faces significant hurdles. Existing medical adaptations (e.g., MedSAM, SAM-Med3D) are typically trained on mixed-organ datasets containing structures like the heart, liver, and lungs. These models struggle to adapt to the brain's unique anatomical characteristics, such as complex tissue organization, indistinct gray-white matter boundaries, and highly diverse lesion morphologies. Furthermore, most SAM-based adaptations convert 3D volumetric scans into 2D slices, neglecting spatial context across slices. Additionally, the reliance on manual interactive prompts limits scalability, as clinicians may struggle to provide accurate prompts for subtle pathologies like epileptic lesions. Finally, the original SAM mask decoder, designed for 2D natural images with clear boundaries, often performs suboptimally on brain MRI data characterized by blurred edges and tissue heterogeneity.

Methodology
The authors propose Brain-SAM, a model tailored specifically for 3D brain MRI lesion segmentation. The architecture extends the SAM2 framework with three primary modifications:

1. 3D Adaptation of the Image Encoder: The original 2D Hiera encoder (a hierarchical Vision Transformer) is adapted to process 3D volumetric data directly. This involves extending four key modules:

   • Patch Embedding: Replacing 2D convolutions with 3D convolutions to partition volumetric inputs into voxel patches.
   • Position Embedding: Extending from 2D to 3D representations to capture spatial relationships across depth, height, and width.
   • Token Sequence: Flattening 3D volumes along all three dimensions to create longer token sequences suitable for volumetric attention.
   • Q-Pooling: Replacing MaxPool2d with MaxPool3d to enable feature downsampling across depth and spatial dimensions.

2. UNETR-Inspired Decoder: To address the limitations of the original SAM decoder in handling fine-grained local textures and boundary details, the authors introduce a decoder inspired by UNETR. This component integrates the global modeling capacity of Transformers with the local fine-grained representation of CNNs via multi-scale feature fusion. Unlike standard UNETR decoders that require resolution transformations, this design leverages the Hiera encoder's inherent hierarchical features for direct fusion and progressive upsampling. The architecture supports a dual-head output: one for fully automatic segmentation and another for generating intermediate features for prompt-based segmentation.

3. Training and Fine-Tuning Strategy:

   • Pretraining: The model is initialized with SAM2 weights and undergoes a pretraining stage on the BraTS2021 dataset (1,000 cases, 4 modalities) exclusively for automatic segmentation to establish robust medical representations.
   • Adapter-Based Fine-Tuning: To adapt to multiple brain lesion datasets (tumors, stroke, epilepsy) while preventing catastrophic forgetting, lightweight adapter modules (bottleneck architecture with 3D convolutions) are inserted before each Hiera block. During fine-tuning, the backbone remains frozen, and only the shared adapter parameters are updated.
   • Prompt Mechanism: Point prompts are used as the primary interaction method, extended to 3D. The mask generation strategy is modified: instead of using updated keys in the dot-product operation, the design fuses processed queries with intermediate feature representations from the UNETR-inspired decoder.
   • Refinement: A post-processing refinement mechanism uses user clicks to locally adjust mask probabilities via 3D Gaussian kernels.

Key Contributions

1. 3D Extension of SAM2: The work successfully extends the SAM2 framework from 2D natural images to 3D medical image analysis, enabling direct volumetric processing without slice-by-slice approximation.
2. Hybrid Decoder Architecture: The introduction of a UNETR-inspired decoder allows the model to support both fully automatic and interactive segmentation within a unified architecture, effectively capturing both local fine-grained details and global contextual dependencies.
3. Comprehensive Dataset Curation: The authors curated a large-scale dataset covering brain tumors, stroke, and epilepsy, enabling the model to support a broad range of clinical applications.

Results
Extensive experiments were conducted on seven public datasets (BraTS-PED, BraTS-SSA, BraTS-MET, BraTS-MEN, ISLES2022, ATLAS v2.0, OPMRI) covering six MRI modalities. Brain-SAM was compared against task-specific models (nnU-Net, UNETR, Swin UNETR) and foundation models (SAM-Med3D, VISTA3D).

• Overall Performance: Brain-SAM achieved superior Dice scores and lower 95% Hausdorff distances (HD95) compared to all baselines. Specifically, compared to nnU-Net, it achieved Dice improvements of 22% on epileptic lesions, 9% on brain metastases, and 6% on meningiomas.
• Small Lesion Segmentation: The model demonstrated significant advantages in segmenting small lesions (e.g., epilepsy), achieving 15%–18% higher Dice scores than strong baseline models. It maintained a better balance between recall and precision for small targets (<5,000 voxels) compared to other methods.
• Ablation Studies:
  • Replacing the 3D-ViT backbone with 3D-Hiera improved Dice scores by up to 14% for small lesions and significantly reduced HD95.
  • Integrating the UNETR-inspired decoder increased the average Dice score by 8% compared to the version without it.
  • The adapter-based fine-tuning strategy proved highly parameter-efficient (39.4M tunable parameters vs. 72.02M for full fine-tuning) while achieving performance comparable to full fine-tuning.

Significance and Claims
The paper claims that Brain-SAM offers a robust pre-trained model for downstream brain MRI analysis tasks. By addressing the specific anatomical and structural challenges of brain tissue—such as indistinct boundaries and small lesion sizes—it provides more reliable and accurate segmentation than existing generalist foundation models. The authors posit that this capability can support better-informed clinical decision-making, particularly in early diagnosis and treatment planning for conditions like epilepsy and brain tumors. The model's ability to function in both automatic and prompt-guided modes aims to reduce manual annotation workloads while maintaining high accuracy. The authors note that the current study is limited to single-modality data and suggest that future work could explore multi-modality integration to further enhance lesion characterization.