[Paper Review] Customized Segment Anything Model for Medical Image Segmentation

Kaidong Zhang, Dong Liu|arXiv (Cornell University)|Apr 26, 2023

Advanced Neural Network Applications121 citations

TL;DR

SAMed customizes the Segment Anything Model (SAM) for medical image segmentation using LoRA finetuning, achieving competitive results with minimal deployment/storage overhead.

ABSTRACT

We propose SAMed, a general solution for medical image segmentation. Different from the previous methods, SAMed is built upon the large-scale image segmentation model, Segment Anything Model (SAM), to explore the new research paradigm of customizing large-scale models for medical image segmentation. SAMed applies the low-rank-based (LoRA) finetuning strategy to the SAM image encoder and finetunes it together with the prompt encoder and the mask decoder on labeled medical image segmentation datasets. We also observe the warmup finetuning strategy and the AdamW optimizer lead SAMed to successful convergence and lower loss. Different from SAM, SAMed could perform semantic segmentation on medical images. Our trained SAMed model achieves 81.88 DSC and 20.64 HD on the Synapse multi-organ segmentation dataset, which is on par with the state-of-the-art methods. We conduct extensive experiments to validate the effectiveness of our design. Since SAMed only updates a small fraction of the SAM parameters, its deployment cost and storage cost are quite marginal in practical usage. The code of SAMed is available at https://github.com/hitachinsk/SAMed.

Motivation & Objective

• Extend SAM to semantic medical image segmentation with meaningful tissue labels.
• Enable efficient fine-tuning by updating a small fraction of SAM parameters.
• Leverage warmup and AdamW to stabilize training and improve convergence.
• Demonstrate competitive performance on the Synapse multi-organ dataset.
• Show that SAMed incurs marginal deployment/storage overhead while remaining SAM-compatible.

Proposed method

• Freeze the SAM image encoder and apply LoRA to its transformer blocks to learn medical features.
• Fine-tune the prompt encoder and the mask decoder (optionally with LoRA) for semantic segmentation.
• Adapt SAM's output to predict k semantic masks corresponding to k classes (including background) and compute final S via Softmax and ArgMax over the class dimension.
• Use cross-entropy and Dice losses with a downsampled ground truth for training supervision.
• Adopt warmup phase and AdamW optimizer to stabilize training and improve convergence.
• Demonstrate compatibility with SAM and show reduced updated parameter count (e.g., 18.81M with LoRA on image encoder only).

Experimental results

Research questions

• RQ1 Can SAM, when fine-tuned with LoRA on medical data, perform semantic segmentation for medical images?
• RQ2 Does updating a small fraction of SAM parameters yield competitive segmentation accuracy with lower deployment/storage overhead?
• RQ3 What training strategies (warmup, AdamW) improve convergence and performance for SAM adaptations on medical data?
• RQ4 How does SAMed compare to state-of-the-art medical segmentation models on Synapse in DSC and HD?
• RQ5 Can SAMed provide meaningful semantic labeling for different tissues while maintaining SAM compatibility?

Key findings

• SAMed achieves 81.88 DSC and 20.64 HD on the Synapse multi-organ dataset, on par with state-of-the-art baselines.
• SAMed attains state-of-the-art performance on pancreas and stomach segmentation within the Synapse results.
• Only a small fraction of SAM parameters are updated (e.g., 18.81M vs 358M original; 5.25% of original size), keeping deployment/storage overhead marginal.
• LoRA applied to the image encoder (and optionally mask decoder) yields better performance than updating only the mask decoder.
• Warmup and AdamW optimizer significantly stabilize training and improve convergence and final loss.
• SAMed remains fully compatible with SAM and can be used as a plug-in to empower SAM for medical image segmentation.