Source-linked AI summary
Customized Segment Anything Model for Medical Image Segmentation
Kaidong Zhang, Dong Liu
TL;DR
Medical image segmentation needs models that handle medical structures and semantic labels without the engineering and storage burden of task-specific networks. SAMed customizes SAM with parameter-efficient finetuning and training strategies, achieving competitive Synapse performance while remaining lightweight to deploy and store.
Problem
Medical segmentation requires domain-specific boundaries and semantic labels, while existing task-specific models demand substantial engineering and deployment resources.
Method
SAMed customizes SAM with LoRA-adapted image encoding plus finetuned prompt encoding and mask decoding, using warmup and AdamW training strategies.
Results
81.88 DSC and 20.64 HD on Synapse are on par with current state-of-the-art baselines after 160 finetuning epochs.
Takeaways & Limitations
SAMed performs medical semantic segmentation while adding only marginal deployment and storage overhead and remaining compatible with SAM.
Abstract
from arXiv · showhide
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.
1 Introduction
SAMed customizes SAM for medical image segmentation, addressing medical-domain boundaries and semantic labels while retaining competitive performance with marginal deployment and storage overhead.
- Medical image segmentation supports computer-aided diagnosis and intelligent clinical surgery by identifying anatomical or pathological tissue structures.
- Transformer-based medical segmentation improves global feature aggregation but typically requires extensive engineering and sizable deployment models.
- Large-scale vision models could unify medical segmentation solutions, but limited medical data and semantic labels hinder direct application.
- SAMed adapts SAM using LoRA in the image encoder while finetuning the prompt encoder and mask decoder for medical semantic segmentation.The method updates only a small fraction of parameters and remains compatible with SAM.
- 81.88 DSC and 20.64 HD on Synapse are on par with current state-of-the-art baselines after 160 finetuning epochs.Warmup and AdamW stabilize finetuning and improve segmentation precision.
2 Related Works
Related work traces medical image segmentation from handcrafted and U-Net-based methods toward Transformer architectures and large-scale pretrained models. It frames efficient finetuning as a way to inject downstream knowledge into such models.
- Early medical segmentation used explicit contour features and Markov random fields, followed by U-Net and variants such as Res-UNet, Dense-UNet, U-Net++, and 3D-Unet.
- Transformers expanded medical image segmentation research and enabled construction of large-scale vision models with billions of parameters.
- Vision models including SAM, SegGPT, and STU-Net attract attention because of their broad zero-shot generalization ability, although limitations remain for downstream use.
- A central research question is how to transfer knowledge from pretrained large-scale models to downstream tasks through improved finetuning strategies.
3 Method
SAMed customizes SAM for automatic semantic segmentation of medical images by freezing its image encoder while adding trainable low-rank updates and finetuning prompt and mask-decoding components. Its method combines LoRA-based adaptation, semantic mask decoding, cross-entropy and Dice supervision, warmup training, and AdamW optimization.
- Overview: SAMed freezes SAM’s image encoder and adds trainable LoRA layers while finetuning the prompt encoder and mask decoder for medical semantic segmentation.The prompt encoder uses a trainable default embedding, enabling automatic segmentation without prompts during inference.
- LoRA in image encoder: LoRA approximates updates to frozen projection layers with a low-rank bypass composed of two trainable linear layers whose rank is much smaller than the input and output dimensions.The bypass uses matrices A and B, with r ≪ min{C_in, C_out}.
- LoRA in image encoder: SAMed applies LoRA to the query and value projection layers of each image-encoder transformer block, which the authors observe performs better than other projection choices.The q, k, v, and o terms denote the corresponding attention projection layers; W_q, W_k, and W_v are frozen while LoRA parameters are trainable.
- Prompt encoder and mask decoder: The mask decoder integrates sparse and dense embeddings with encoded image features, predicts semantic masks for each class, and postprocesses them into the final segmentation result.SAMed modifies the segmentation head to produce k semantic masks corresponding to the predefined semantic labels.
- Training strategies: SAMed supervises finetuning with a weighted combination of cross-entropy and Dice losses after downsampling the ground truth to the output resolution.The weights λ1 and λ2 balance the influence of the two loss terms.
- Training strategies: Warmup stabilizes early training, followed by exponential learning-rate decay, while AdamW improves performance compared with SGD in SAMed.The authors attribute the AdamW observation to consistency between the optimization recipes for SAM and SAMed.
4 Experiments
Experiments on Synapse evaluate SAMed against established methods and test its customization and training choices. Results support competitive segmentation with reduced model overhead, while warmup and AdamW improve convergence.
- Quantitative comparison: SAMed achieves highly competitive performance on Synapse but is not the best among the compared state-of-the-art methods.The comparison includes U-Net, Att-UNet, TransUnet, SwinUnet, MissFormer, TransDeepLab, HiFormer, and DAE-Former.
- Quantitative comparison: SAMed achieves state-of-the-art performance for pancreas and stomach segmentation while adding only marginal deployment and storage overhead.The method customizes only a small fraction of SAM parameters and remains compatible with the segment anything system.
- Qualitative comparison: Qualitatively, SAMed produces smoother and more correct segmentation regions than TransUnet, SwinUnet, and DAE-Former.The authors attribute this to SAM’s feature extraction ability and SAMed’s finetuning strategies.
- Ablation study: Finetuning the image encoder with LoRA improves segmentation accuracy over finetuning only the segmentation head.The authors associate this improvement with extracting meaningful medical-image features before mask decoding.
- Ablation study: SAMed_s reduces deployment and storage overhead but performs worse than SAMed when LoRA is applied to the mask-decoder transformer instead of updating all decoder parameters.SAMed_s may substitute for SAMed under strict deployment or storage limits.
- Training strategies: Warmup stabilizes training and AdamW further lowers the later-stage loss, producing a more stable process that converges to a lower loss value.The baseline omits warmup and uses SGD instead of AdamW.
5 Conclusion
The paper demonstrates that customizing SAM can support competitive medical image segmentation on Synapse. SAMed combines LoRA-based image-encoder adaptation with suitable training strategies while retaining marginal deployment and storage overhead.
- Conclusion: SAMed demonstrates that a customized large-scale model can achieve highly competitive results against well-designed medical image segmentation methods.The conclusion frames this as evidence for the feasibility of customizing large-scale models for medical image segmentation.
- Conclusion: SAMed uses a LoRA-adapted image encoder and suitable training strategies to achieve competitive Synapse results with marginal deployment and storage overhead.The method is described as fully compatible with SAM for practical usage.