Attempting to Understand Scientific Subject Matter (feat. MedSAM)

Following on from previous posts, "Segment Anything in Medical Images" will reprise its starring role in this month's blog entry - our topic now turning from writing style to subject matter. Specifically, the critical reading & understanding of a scholarly work's content and its value in the context of the relevant sub-field.

Context & Problem

Foundation Models (FMs) are notably large (i.e., having many parameters), deep (i.e., having many layers) neural networks trained on substantial amounts of data so they can be used across many, varied tasks. This general-purpose applicability constitutes their ⭐ grand appeal ⭐. A combination of architecture, training dataset properties, and training protocol means FMs can be put to work for a range of tasks out-of-the-box or with minimal fine-tuning instead of having to develop, train, and deploy an independent model for each use case.

One such FM - Meta AI's Segment Anything (SAM) model - shook the sub-field of image segmentation when released in 2023 for its interactivity, adaptability, and efficacy in segmenting all kinds of...stuff ("anything!") from images. SAM takes various prompts as input - like clicks or bounding boxes - and segments the referenced object/s through pixel-wise classification. Its ability to conduct this segmentation on unfamiliar objects and images means the model has learned a strong enough notion of what different structures are that you may hear SAM and "zero-shot generalisation/performance" 🎯 in the same sentence often; perhaps with adjectives like "revolutionary", "ground-breaking", and/or "transformative" in the mix.

Taking the sales pitch on FMs at its word, SAM's introduction implied no more hefty NVIDIA invoices for training that specialised car/fish/meteor/etc. segmentation network you always wanted for Christmas. SAM could do it all in one! 🎁 Everyone gets a happy ending, a tax deduction, and I can take a nap instead of writing the rest of this blog! 💤 Just kidding. Enter our complex, misunderstood anti-hero - medical imagery:

Figure 1: Visual differences in medical and natural images.

Because of the marked difference in features learned between image types (Figure 1), leveraging a general segmentation model trained on natural images for the same task on medical images yields sub-optimal performance. That is to say, the "meaning" learned from extracted features - like edges/boundaries, intensity distributions, and textures - tends to vary greatly between a natural and medical image. Intuitively, outlining cortical regions in a neonatal brain MRI requires different know-how to picking a cow out of a field 🐄. So, even an FM as revolutionary, ground-breaking, and transformative as SAM was observed falling short compared to traditional, medical task-specific deep learning models, especially where target objects had weak boundaries (e.g., low contrast between different tissues).

Main Idea: A Hybrid Solution

Building on the zero-shot generalisation capability of SAM with domain knowledge learned from further training on medical images to produce one universal model that can rival task-specific network performance and be applied to a wide variety of medical structures and scan types/modalities is the core contribution of this seed paper. The paper proposes said model as MedSAM. 

Figure 2: Key MedSAM components, including medical domain data (blue) and foundation model/SAM architecture (orange)

MedSAM retains SAM's core architecture and pre-trained weights to preserve useful functionality - like how to interpret a bounding box prompt and general object segmentation knowledge. The paper's novelty extends that knowledge to medical imagery by fine-tuning SAM on a massive dataset of 1.5M image-mask pairs spanning 10 modalities (e.g., MRI, CT, X-ray, ultrasound), >30 cancer types, and multiple imaging protocols, where each mask was manually labelled by a human expert.

Figure 2 captures the gist of it, where existing SAM components are outlined in orange and new components (the data!) are in blue.

Data Properties

MedSAM's fine-tuning success is mostly underpinned by the diversity and size of the dataset, enabling comprehensive learning of medical image features and their semantics. Including 10 modalities and a variety of pathological states allows MedSAM to learn a rich representation of medical structures, accounting for unique characteristics across imaging types and anatomical conditions  (e.g., what a particular organ may look like in an MRI and a CT volume, healthy, damaged, underdeveloped, and the like 🤕).

The authors collected, combined, transformed/preprocessed, and standardised medical images from a multitude of sources to ensure consistency, compatibility, quality, in the hopes of improving performance. Data-centric prep involved:

• Image format conversions - for singular image handling & model compatibility.
• Intensity value normalisation - so all intensities fall within the same range, and outlier values do not destabilise learned representations.  
• Patch extraction, resizing, and interpolation - so all images had the same dimensions for uniform input processing without unintended artifacts/deformities (e.g., stretched objects) that inhibited training ("confusing the model"). 
• Bounding box simulation for masks using expert annotations as input with a random perturbation of 0 - 20 pixels to factor in the variation you'd expect from actual, user-specified boxes. 
  

Model Architecture

In terms of other notable design choices, the paper authors stayed true to their goal of producing a universal solution by opting for a 2D model that could be adapted to 3D if and when required. Rather than multiple prompt types, MedSAM training focused on bounding boxes for precision between clinical applications that demand varying granularity levels (e.g., using a single point to segment a liver tumour may result in segmenting the entire liver, so it's a safer bet to outline just the tumour with a box).

As outlined in the original SAM proposal, the core model architecture includes:

• A vision transformer (ViT)-based image encoder 🖼️ that maps the input medical image to a high-dimensional embedding space (feature extraction).
  
• A prompt encoder 📦 that transforms bounding box prompts to prompt embeddings using positional encoding (capturing where the box is in the image, spatially). This was fixed during fine-tuning as bounding box encoding did not require adaptation.
• A mask decoder 🎭 that combines and processes the image and prompt embeddings using cross-attention to output a segmentation mask (the input image with pixels now classified).
  

Overall, the ViT model underlying MedSAM has 12 transformer blocks, where each block contains a multi-head self-attention component for capturing context (relationships between patches), a multi-layer perceptron component for feature refinement, and a layer normalisation component for training stability (makin' the math work good). The two transformer layers that make up the mask decoder have additional convolutional layers to increase the combined embedding's resolution before interpolation ensures the output mask matches the original input dimensions.

Training Protocol/Configuration

The MedSAM paper lists but doesn't elaborate on much of the tuning process or decisions for settings used to train the final model - including using the AdamW optimiser for parameter updates, a learning rate of 1e-4 to control step size during optimisation, weight decay of 0.01 to control network complexity (reduce overfitting), a batch size of 160, and no data augmentation. Parameter adjustment is mentioned earlier but detail such as search space/ranges and approaches are not specified. Such omissions may be common in ML academia due to publication pace or lab constraints....I'm not an experienced enough scholar to say! But it would've been cool to hear how the pros do it. 🥲

In contrast, the unweighted sum between cross-entropy loss and dice loss was used because of its proven efficacy in a wide range of medical image segmentation tasks. Both a region-based metric (Dice) and boundary-based metric (Normalised Surface Distance) were used to maximise evaluation coverage.

After a model checkpoint was produced following 150 training epochs, MedSAM was evaluated on 146 validation tasks without further changes or fine-tuning.

Results

Experiments were run across a total of 86 internal validation tasks, i.e., using train/val/tuning subsets from the same distribution and 60 external validation tasks, i.e., out-of-distribution test data from different scanners and patients than those that appeared during training.

Evaluation pitted MedSAM against the base SAM model and two leading segmentation architectures - nnU-Net and DeepLabV3+ - trained per modality to represent task-specific performance benchmarks.

Quantitative results showed MedSAM versatility and performance consistently surpassing respective competitors:

• MedSAM ranked first on most internal validation tasks with a narrower Dice Similarity Coefficient (which measures the overlap between predicted and ground truth segmentation masks) distribution than its competitors, illustrating greater performance than task-specific models and FM-like robustness across different tasks.
  
  
• SAM performance was lowest across most internal validation tasks, except some RGB image tasks - speculated due to better boundary distinctions that would be easier for SAM to identify as per its training (natural image feature similarity).
  
  
• The performance margin between MedSAM and specialist models only grew in external validation tasks, highlighting the limited generalisability of the latter when faced with unseen data. SAM performance remained similarly low.
  
  
• This trend continued when experimenting with entire, unseen modalities and tasks. MedSAM surpassed the three other models with improvements of up to 10%. 🤯

To sum, MedSAM's ability to out-perform the other models during external validation, specifically, was a strong testament to how a combination of foundation model architecture & pre-training could be bolstered using a broad enough dataset to achieve promising results for universal model application in real-world clinical settings.

tl;dr huge success! 😎💚

Conclusion & Follow-ups

The MedSAM model & paper were a big step forward for research surrounding FMs leveraged in the medical domain and keep in mind - the field is still moving fast! Meta AI released a second version of SAM, aptly named Segment Anything 2 (SAM 2) (Figure 3), in August 2024 with multiple architectural updates, including a memory module that strengthens 3D and video segmentation capability by providing a mechanism for the network to learn from past predictions. For example, learning that part of an ulcer is still an ulcer, even when the rest of it is obscured by fluid in some endoscopy frames. 🩸📽️

Figure 3: SAM 2 architecture

The authors of MedSAM recently made record time applying their adaptation methodology to SAM 2 to produce (you'll never guess the name), MedSAM 2! The successor was released last month (April 2025) and, although the paper is still in pre-print, the results are just as exceptional as the original.

The proposal of MedSAM 2 (Figure 4) comes with an empirical assessment of the second model's performance on 3D imagery (compared to only the promise of 3D adaptation from the first) and video. The dataset is comprised of >450k 3D image-mask pairs and 76k annotated video frames, and the domain adaptation methodology was adjusted to account for the move away from the first iteration's 2D-centric design. That being said, the enhancements for video are still backwards-compatible with 3D & 2D processing, meaning new features could be harnessed across all use case categories.

Figure 4: MedSAM 2 prediction process (bi-directional, slice-wise propagation of segmentation masks)

Between both models, there is still a question of how modality imbalances in the training data (e.g., more CT & MRI samples than other scan types) and task variety (e.g., MedSAM 2 focuses on organs & lesions and neither paper includes results on musculoskeletal structures) could be explored to either further improve performance in individual tasks or broaden the range of possible applications.

In MedSAM 2 specifically, prompting remains limited to bounding boxes and inference still requires GPU computation, meaning future research directions also emphasised practical optimisations alongside model performance to push along adoption at your local hospital ASAP! Super interesting territory (that I should hopefully get to poke & prod at while struggling through this semester's project)!! Fun! 🙂✨🩻🔎