Embedded framework for clinical medical image segment anything in resource limited healthcare regions

Abstract

The emergence of portable imaging devices improves medical image acquisition efficiency in resource-limited regions, but a shortage of medical personnel still limits timely diagnosis. We propose Embed-MedSAM, a fully automatic segmentation model with low deployment cost. Built on MedSAM, it integrates a lightweight RepViT encoder to reduce computation and applies two-stage distillation on over one million multimodal medical images to preserve the original model’s visual representation. A self-prompting mechanism is also introduced, where the model generates pseudo labels to guide fine-grained segmentation. The training jointly optimizes KL divergence and segmentation losses to improve accuracy under prompt-free conditions. Embed-MedSAM shows excellent performance on 17 benchmark datasets covering 7 imaging modalities. Without external prompts, it improves average Dice score by nearly 16% over the second-best model. It also runs at nearly 30 FPS on iPhone 14, showing strong potential for real-world deployment.

Introduction

In resource-constrained environments such as conflict zones, remote areas, and economically disadvantaged regions, the lack of medical resources remains a major barrier to improving healthcare delivery^1,2. In recent years, the development and deployment of portable clinical imaging devices^3–6 have significantly improved access to healthcare in these areas. These devices allow local facilities to acquire high-quality medical images more easily.

However, relying on imaging equipment alone does not solve the core problem. These regions often suffer from a severe shortage of skilled professionals, especially those trained to interpret complex medical images. In such settings, the ability to quickly analyze images and detect lesions is critical. Delays in diagnosis may result in missed treatment opportunities. Therefore, efficiently segmenting and diagnosing medical images across different modalities, under hardware-limited conditions, remains an urgent and unsolved challenge.

To address this need, computer vision-based automated segmentation techniques have been introduced in medical image analysis. These methods aim to identify and delineate regions of interest (ROIs), such as organs and tissues⁷. Recently, automatic segmentation approaches have gained attention for improving both consistency and precision^8,9. At the same time, advances in deep learning¹⁰ have enabled the widespread adoption of neural networks for medical image segmentation. These include traditional Convolutional Neural Networks (CNNs)¹¹ and, more recently, Vision Transformers (ViTs)¹².

U-shaped architectures based on CNNs model the relationship between low-level and high-level semantic features for mask prediction^13–15. These methods can produce accurate segmentation results in known domains. However, their performance often degrades in unseen domains.

To tackle the challenge of Domain Generalization (DG), existing approaches adopt strategies such as multi-task learning, data augmentation, and domain synthesis. These techniques aim to increase the diversity of learned feature representations^16–18. Nevertheless, due to limited model capacity, many of these task-specific models must be retrained from scratch when applied to different imaging modalities. This creates a major obstacle in building segmentation models with strong generalization and versatility across diverse medical domains.

Vision Transformers (ViTs)¹² address this limitation by capturing long-range dependencies through self-attention. When trained on large-scale datasets, ViTs benefit from greater model capacity and reduced inductive bias. Recently, the Segment Anything Model (SAM)¹⁹ has achieved major breakthroughs in image segmentation. SAM has demonstrated strong generalization in natural images and performs well across various scenarios through prompt-based interaction. Leveraging its powerful pre-trained encoder, several studies^20–23 have explored SAM’s potential in general medical image segmentation.

However, although MedSAM (as shown in Fig. 1) is a model specifically designed for medical image segmentation within the current SAM series, it still faces two major challenges that remain to be addressed in practical applications.

Fig. 1. Overview of the MedSAM framework.

The input image is first processed by a mask encoder to extract image features. These features are combined with prompt information, such as masks, points, boxes, or text. The prompt encoder processes these inputs. Then, the mask decoder uses both the image and prompt embeddings to generate the final segmentation output.

(1) MedSAM and its variants have extremely large parameter sizes, typically exceeding 600 million, which results in high computational costs during inference. Two common approaches have been proposed to address this issue: (1) equipping the algorithm with dedicated high-performance computing devices; or (2) deploying the model on the cloud and transmitting inference results remotely. However, both solutions are impractical in regions with limited hardware and network resources-such as conflict zones, remote mountainous areas, or high-altitude regions-and the latter also raises concerns about data privacy and security.

(2) MedSAM’s inference process relies heavily on user-provided prompts (e.g., points or bounding boxes) to guide the generation of segmentation masks. Although bounding box-based prompts can yield relatively accurate results, users are typically required to roughly indicate the location of the lesion or target structure in the image. This imposes a certain level of medical expertise on the user, which limits the model’s applicability in underdeveloped or non-professional environments.

To address the above challenges, we propose Embed-MedSAM (Embedded Medical Image Segment Anything Model), as shown in Fig. 2. Specifically, Embed-MedSAM is an optimized extension of MedSAM²¹, with improvements made through a two-stage process:

Fig. 2. Overall technical framework of the Embed-MedSAM model.

a Image encoder distillation stage. a shows the replacement of the original MedSAM encoder with a lightweight RepViT encoder. Step 1 illustrates the masked image reconstruction pretraining on ImageNet-21K, using reconstruction loss to transfer general visual knowledge. Step 2 shows the knowledge distillation on the COSMOS 1050K medical dataset. It uses MSE losses to guide the lightweight encoder to align with the semantic features of MedSAM. b Self Prompt Training Framework. This stage aims to remove the need for manual prompts. The model generates a self-guided prompt from the initial prediction, encodes it using the prompt encoder, and feeds it into the frozen decoder. To suppress the accumulation of pseudo-prompt errors, we introduce a KL divergence loss during training to enforce semantic consistency between the two predictions. In addition, an L1 loss between the pseudo-prompt box and the ground truth mask box is used to enhance geometric alignment. At the mask level, Dice loss and BCE loss are combined to optimize the overlap and pixel-wise consistency between the final prediction and the ground truth. c Speed-accuracy trade-off analysis (for reference only). c compares the FPS and Dice scores of Embed-MedSAM across multiple platforms and datasets. It shows that the model achieves a good balance between inference speed and segmentation performance.

(1) We replace the computationally expensive image encoder in MedSAM with a more efficient lightweight RepViT model²⁴ (as shown in Fig. 3) to significantly reduce computational cost and inference latency. To preserve the strong visual representation capability of MedSAM, we design a two-step distillation strategy: first, we perform masked image pretraining on the general-domain dataset ImageNet-21K, combined with reconstruction supervision from the original MedSAM to transfer general visual knowledge; second, we introduce a optimization of Mean Squared Error (MSE) loss on the large-scale medical imaging dataset COSMOS 1050K, guiding the lightweight encoder to align with the semantic features of MedSAM and thereby effectively inheriting its domain adaptability in medical image tasks.

Fig. 3. Architecture of light-weight encoder (image encoder, RepViT [24]).

The architecture consists of a stem layer and four main stages. Each stage operates at a different resolution: H/4 × W/4, H/8 × W/8, H/16 × W/16, and H/32 × W/32, where H and W are the height and width of the input image. C_i represents the number of channels in the i-th stage, and B is the batch size. Each stage contains RepViT blocks, which are based on depthwise separable convolution (DWConv) and feed-forward networks (FFNs). Downsampling is applied between stages. The model also integrates lightweight SE attention and adapter modules to enhance feature representation with low computational cost. The bottom part of the figure shows the internal structure of each block, including 3 × 3 depthwise convolution, 1 × 1 depthwise convolution, and FFN modules. These components are used in both training and inference.

(2) We propose a self-prompting training framework for medical image segmentation without manual prompts. The method uses the model’s initial forward prediction to generate a pseudo mask, which is converted into a bounding box using the minimum enclosing rectangle. This pseudo box is then refined by the prompt encoder to guide a second, more accurate prediction. To reduce accumulated errors from pseudo prompts, we introduce a KL divergence loss to enforce consistency between the two predictions. An additional L1 loss is applied between the pseudo box and the ground truth box to improve geometric alignment. At the mask level, Dice loss and binary cross-entropy loss are used to optimize overlap and pixel-wise accuracy. This framework requires no human-provided prompts during training or inference, enabling better generalization and deployment in real-world scenarios.

Results

Training details of image encoder distillation

Step 1: Masked image pretraining. In this step, we integrate RepViT as the image encoder within the MedSAM framework, while keeping the prompt encoder and mask decoder unchanged. We conduct masked image pretraining on 0.5% of the ImageNet-21K dataset without using any label annotations. During pretraining, the reconstruction targets are the feature maps produced by the original MedSAM image encoder. The model is optimized using a Mean Squared Error (MSE) loss. The training is performed for 10 epochs with a batch size of 16, using the AdamW optimizer with a learning rate of 2e-3, β₁ = 0.9, and β₂ = 0.95. The training is conducted on 4 NVIDIA A800 80GB GPUs.

Step 2: Knowledge distillation training. In this step, we distill semantic features from the MedSAM image encoder into the lightweight encoder pre-trained in the previous stage. This training is performed on the COSMOS 1050K dataset²⁰. The model is trained for 100 epochs with a batch size of 16, using the Adam optimizer with an initial learning rate of 0.0001 and a weight decay of 0.001. To prevent overfitting, we apply early stopping, terminating training when the validation loss no longer decreases. Additionally, a learning rate scheduler reduces the learning rate by a factor of 10 once the validation loss plateaus. The training is conducted on 4 NVIDIA A800 80GB GPUs.

Training details of self-prompting training framework

At this stage, training is conducted on the source domain datasets listed in Table 1, using their training and validation sets for model optimization. Evaluation is performed on both the test sets of the source domains and the unseen external datasets. By default, the AdamW optimizer is used, with an initial learning rate of 1e-3 and a weight decay of 0.01. The learning rate is dynamically adjusted based on validation loss using the ReduceLROnPlateau strategy. The model is trained for 100 epochs with a batch size of 8. All experiments can be completed using two NVIDIA 2080Ti (11GB) GPUs.

Table 1.

Details of the proprietary datasets used in the experiments

Modality	No.	Dataset	Resolution	Samples
	SOURCE DOMAINS (Internal datasets)
Dermoscopy	${\mathcal{S}}^1$	ISIC-201834,35	Variable	3694
X-ray	${\mathcal{S}}^2$	PCXA36,37	Variable	704
	${\mathcal{S}}^3$	DRIVE38	584 × 565	40
Fundus	${\mathcal{S}}^4$	REFUGE2-Disc39	Variable	1200
	${\mathcal{S}}^5$	REFUGE2-Cup39	Variable	1200
Colonoscopy	${\mathcal{S}}^6$	CVC-ClinicDB40	384 × 288	612
Ultrasound	${\mathcal{S}}^7$	UDIAT41	Variable	163
	${\mathcal{S}}^8$	CCH-TRUSPS42,43	Variable	2879
Microscopy	${\mathcal{S}}^9$	DSB-201844	Variable	670
MRI	${\mathcal{S}}^{10}$	BraTs 202145	240 × 240 × 155	2000

	UNSEEN DOMAINS (External datasets)
Dermoscopy	${\mathcal{T}}^1$	PH246	767 × 576	200
X-ray	${\mathcal{T}}^2$	NIHCXR47	512 × 512	100
Fundus	${\mathcal{T}}^3$	STARE48	700 × 605	20
Colonoscopy	${\mathcal{T}}^6$	CVC-ColonDB49	574 × 500	380
Ultrasound	${\mathcal{T}}^7$	BUSI50	Variable	780
Microscopy	${\mathcal{T}}^9$	TNBC51	512 × 512	50

Evaluation metrics

To comprehensively evaluate medical image segmentation performance, we employ three standard metrics: Dice score, Hausdorff Distance (HD), and Intersection over Union (IoU). These metrics quantify the similarity between predicted masks and the ground truth (GT). Additionally, we report model parameters, Floating Point Operations (FLOPs), and Frames Per Second (FPS) to assess computational efficiency.

It is important to note that, in comparative experiments, all competing models report only their average best Dice and HD scores [25] on the test sets. In contrast, the proposed Embed-MedSAM model reports the mean standard deviation over five runs with different random seeds, or the mean value alone in specific cases, to more accurately reflect model stability and overall performance. Construction of subsequent experiments reference [25] and [26].

Dataset

In the first stage of training, we used the ImageNet-21K dataset for masked image pretraining and the COSMOS 1050K medical imaging dataset for knowledge distillation. It is important to note that COSMOS 1050K consists of multiple public medical datasets. Some of these datasets require official authorization before use. Therefore, in this study, we only used two types of datasets: (1) publicly available datasets for which we successfully obtained usage permission, and (2) datasets that are freely available under their original licenses without the need for extra authorization. We strictly followed all terms of use and ethical requirements. The data was used solely for research purposes, with no commercial use or policy violations involved.

In addition, some public datasets in COSMOS 1050K may overlap with the “seen" or “unseen" datasets listed in Table 1. To ensure fair evaluation and avoid potential data leakage, we applied two key measures: 1. During the training of Embed-MedSAM, we only used the training and validation sets from the sub-datasets in COSMOS 1050K. All test sets were strictly excluded. 2. We used a filtered subset of COSMOS 1050K. Our team manually checked each unseen dataset against the COSMOS 1050K list. Any overlapping or potentially overlapping datasets were completely removed from the training data. With these two measures, we effectively avoided any overlap between COSMOS 1050K and the test datasets used for both internal and external evaluation. This ensures the objectivity and reliability of the model’s evaluation results.

In the second training stage, we introduce 16 downstream task datasets listed in Table 1. These datasets are divided into two categories: SOURCE DOMAINS (internal datasets) and UNSEEN DOMAINS (external datasets). The internal datasets are split into training, validation, and test sets based on their official splits. The training and validation sets are used to train both the baseline models and our proposed Embed-MedSAM. The external datasets are used only for inference. They are not involved in any training and are used to evaluate the model’s generalization ability on unseen domains.

Supervised segmentation performance evaluation

To validate the effectiveness of the proposed lightweight Embed-MedSAM framework in supervised medical semantic segmentation tasks, we conducted systematic comparisons with mainstream fully automatic segmentation architectures and several lightweight SAM variants. These comparative experiments were performed across six medical imaging modalities, as listed in Table 1. It is important to note that models such as the original SAM and MedSAM, despite their strong generalization capabilities, are excluded from this comparison. Their extremely large parameter sizes-over 60 times larger than those of lightweight SAM models-make them impractical for deployment on resource-constrained devices.

Regarding the training strategy, all lightweight SAM models that support prompt inputs were further trained in a supervised manner on each training dataset, starting from their officially released pretrained weights. Training was conducted using the single-point prompt mode, where each prompt point was generated by computing the centroid of each instance using the ConnectedComponentsWithStats function in OpenCV. In addition, these models also support an automatic mask generation mode (denoted by ✗), which was used to evaluate their performance under prompt-free conditions.

The training strategies for all compared models-including network architectures, hyperparameter settings, and training procedures-strictly followed the descriptions provided in their original publications. The datasets used for evaluation include ${\mathcal{S}}^1,{\mathcal{S}}^2,{\mathcal{S}}^3,{\mathcal{S}}^6,{\mathcal{S}}^7$, and ${\mathcal{S}}^9$ from Table 1. Each dataset was split into training, validation, and test sets according to its official protocol. The final results on the test sets after supervised training are reported in Table 2.

Table 2.

Comparison of Embed-MedSAM with other state-of-the-art frameworks in universal medical image segmentation (source domains, internal datasets)

	Manual	Dermoscopy ${\mathcal{S}}^1$		X-ray ${\mathcal{S}}^2$		Fundus ${\mathcal{S}}^3$		Colonoscopy ${\mathcal{S}}^6$		Ultrasound ${\mathcal{S}}^7$		Microscopy ${\mathcal{S}}^9$		Efficiency
Method	Prompt	Dice ↑	HD ↓	Dice ↑	HD ↓	Dice ↑	HD ↓	Dice ↑	HD ↓	Dice ↑	HD ↓	Dice ↑	HD ↓	Params ↓	FLPOs ↓
Subtable A. With mainstream fully automatic segmentation architectures
U-Net13	✗	82.87a	180.90a	93.85a	103.51a	79.13a	68.22a	83.91a	130.96a	69.24a	131.62a	88.16a	130.28a	13.40M	497.91G
U-Net++52		82.69a	175.04a	95.31a	75.84a	80.61a	65.03a	85.77a	152.96a	72.50a	137.84a	90.48a	112.12a	9.16M	558.46G
Att-UNet14		83.97a	170.78a	95.80a	66.16a	80.72a	66.48a	86.90a	156.72a	71.02a	107.70a	91.12a	113.64a	8.72M	269.07G
nnUNet7		84.96a	126.19a	96.06a	68.24a	81.71a	64.16	88.38	127.28a	75.22a	119.66a	91.61	121.16a	34.29M	554.40G
H2Former15		82.12a	191.39a	95.54a	75.44a	81.46a	65.09a	84.66a	142.08a	70.30a	110.98a	90.17a	117.69a	33.71M	536.96G
TransUNet53		84.28a	134.80a	96.27a	56.85a	81.68a	64.91	86.00a	151.14a	71.07a	123.80a	90.03a	109.56a	109.54M	865.31G
ADS18		84.14a	172.84a	94.89a	86.68a	80.48a	68.49a	87.70a	117.74a	72.55a	136.08a	90.32a	115.24a	14.07M	503.59G
CIAug17		83.91a	141.07a	95.82a	73.18a	80.45a	65.72a	87.69	106.68	71.78a	134.84a	90.58a	113.72a	14.26M	500.23G
MADGNet16		85.02a	131.84a	96.22a	67.41a	81.89	64.73a	88.20a	107.16a	72.75a	131.24a	91.38a	98.04a	31.40M	210.65G
Embed-MedSAM (Mean)		88.49	99.30	97.30	52.14	81.69	75.32	93.29	60.38	85.76	80.37	91.29	96.74	10.6M	33.5G
(Standard Deviation)		± 0.23	± 3.45	± 0.30	± 1.76	± 0.54	± 1.29	± 0.17	± 2.27	± 0.46	± 2.79	± 0.43	± 1.44	-	-
Subtable B. With prompt-free lightweight SAM series models
MobileSAM54	✗	86.19a	168.48a	94.75a	175.50a	25.56a	254.24a	77.29a	284.32a	61.51a	342.22a	61.49a	307.44a	9.79M	39.71G
RepViT-SAM55		85.73a	157.66a	94.27a	170.27a	25.10a	263.76a	78.75a	296.36a	61.62a	440.75a	60.54a	287.63a	9.57M	23.64G
EdgeSAM56		85.86a	153.97a	94.52a	183.37a	25.62a	257.85a	76.41a	328.79a	60.24a	400.69a	59.68a	357.20a	9.60M	22.10G
EfficientSAM-S33		87.02a	162.45a	94.89a	169.20a	25.84a	261.51a	78.54a	305.36a	62.05a	430.18a	60.56a	336.55a	25.38M	32.51G
EfficientViT-SAM57		87.18a	151.29a	96.13a	128.60a	25.95a	252.39a	82.48a	317.51a	65.44a	354.98a	66.32a	292.69a	34.80M	89.10G
SAM-Lightening58		86.99a	165.80a	95.83a	115.19a	25.52a	261.25a	80.14a	292.33a	62.59a	376.55a	62.68a	306.66a	19.26M	52.46G
ESP-MedSAM25		88.52	92.42	96.83	40.68	82.42	62.64	92.93	56.32	85.28	82.32	92.24	85.93	28.46M	55.88G
Embed-MedSAM (Mean)		88.49	99.30	97.30	52.14	81.69	75.32	93.29	60.38	85.76	80.37	91.29	96.74	10.6M	33.5G
(Standard Deviation)		± 0.23	± 3.45	± 0.30	± 1.76	± 0.54	± 1.29	± 0.17	± 2.27	± 0.46	± 2.79	± 0.43	± 1.44	-	-
Subtable C. With no external prompt lightweight SAM series models
MobileSAM54	Point	87.97a	105.34a	96.25a	50.94a	69.31a	94.41a	81.83a	86.46a	66.48a	107.78a	87.42a	131.00a	9.79M	39.71G
RepViT-SAM55		88.00a	106.75a	96.07a	52.71a	67.76a	98.66a	81.81a	154.65a	68.38a	103.09a	88.81a	127.17a	9.57M	23.64G
EdgeSAM56		88.10a	100.92a	96.18a	51.60a	68.04a	92.40a	81.76a	105.06a	67.64a	105.51a	87.35a	113.46a	9.60M	22.10G
EfficientSAM-S33		87.11a	108.12a	96.40a	49.77a	76.32a	79.41a	82.81a	96.35a	71.17a	113.57a	88.41a	129.08a	25.38M	32.51G
EfficientViT-SAM57		88.49a	103.61a	96.43a	49.07a	78.16a	77.16a	85.16a	102.72a	74.71a	113.18a	89.37a	116.42a	34.80M	89.10G
SAM-Lightening58		88.28a	101.64a	96.41a	50.35a	74.84a	87.32a	83.70a	97.79a	73.07a	129.53a	89.18a	111.67a	19.26M	52.46G
Embed-MedSAM (Mean)	✗	88.49	99.30	97.30	52.14	81.69	75.32	93.29	60.38	85.76	80.37	91.29	96.74	10.6M	33.5G
(Standard Deviation)		± 0.23	± 3.45	± 0.30	± 1.76	± 0.54	± 1.29	± 0.17	± 2.27	± 0.46	± 2.79	± 0.43	± 1.44	-	-

Bold indicates the best result, and underline indicates the second-best result.

Universal medical image segmentation refers to the model’s ability to maintain consistent performance across different imaging modalities and anatomical structures. This capability is particularly important in real-world scenarios where medical imaging data are highly heterogeneous.

^aScores marked with “^a” indicate that Embed-MedSAM has a statistically significant advantage (p < 0.05) over the specialist models.

Mainstream fully automatic segmentation architectures. As shown in the table, although mainstream fully automatic segmentation architectures are comparable to the lightweight SAM variants in terms of parameter size, their inference cost-measured by FLOPs-is significantly higher. For example, MADGNet requires 210.65 GFLOPs, whereas the proposed model consumes only 33.5 GFLOPs, making the former 6 times more computationally expensive. This highlights the inefficiency of traditional architectures for practical deployment, especially on resource-limited platforms. In terms of segmentation accuracy, the proposed model achieves the best results on the ${\mathcal{S}}^1,{\mathcal{S}}^2,{\mathcal{S}}^6$, and ${\mathcal{S}}^7$ datasets, and ranks close to second-best on ${\mathcal{S}}^3$ and ${\mathcal{S}}^9$. Overall, compared to existing mainstream fully automatic segmentation methods, the proposed Embed-MedSAM model offers the best trade-off between lightweight design and segmentation accuracy. This makes it particularly suitable for deployment in medical image segmentation tasks on resource-constrained edge devices or mobile platforms.

Prompt-free lightweight SAM series models. Meanwhile, we conducted a comprehensive comparison between Embed-MedSAM and other lightweight SAM variants. Under the no-prompt setting, the proposed model significantly outperforms all competing models-except ESP-MedSAM-across all datasets in terms of segmentation accuracy. Although ESP-MedSAM achieves slightly higher accuracy than our model on most datasets, the performance gap is marginal. For instance, on the ${\mathcal{S}}^1$ dataset, ESP-MedSAM achieves 88.52%, which is only 0.13% higher than our model (88.39%). Similarly, on the ${\mathcal{S}}^2$ dataset, the difference is just 0.11%. However, in terms of model efficiency, our method offers substantial advantages. ESP-MedSAM has 28.46M parameters-nearly three times more than our model (10.6M). Its inference complexity is also significantly higher, with 55.88 GFLOPs compared to only 33.5 GFLOPs for our model.

In addition, as shown in Table 8, we further validate the lightweight advantage of Embed-MedSAM. Although ESP-MedSAM and Embed-MedSAM achieve comparable segmentation accuracy, ESP-MedSAM cannot run on small embedded devices such as the iPhone 14 due to its large model size and high inference complexity. In contrast, Embed-MedSAM can be successfully deployed on the iPhone 14, achieving a real-time inference speed of nearly 30 FPS. These results demonstrate that Embed-MedSAM not only maintains high segmentation accuracy but also offers excellent deployment efficiency, making it well-suited for medical image segmentation tasks on mobile and resource-constrained devices.

Table 8.

Comparison of Embed-MedSAM and other SAM series models in terms of unseen domain generalization and inference efficiency on the iPhone 14 embedded platform

Row	Dice (Avg.) ↑	HD (Avg.) ↓	Params ↓	FLOPs ↓	FPS ↑
SAM19	Out of memory
MedSAM21	Out of memory
ESP-MedSAM25	Out of memory
EfficientSAM-S33	Out of memory
EfficientSAM-Ti33	60.29	196.62	9.80M	38.20G	4.90
MobileSAM54	53.54	258.24	9.79M	39.71G	4.90
RepViT-SAM55	51.84	248.12	9.57M	23.64G	37.67
EdgeSAM56	52.04	266.04	9.60M	22.10G	38.70
EdgeSAM-RPN56	54.14	249.04	9.80M	22.20G	34.10
Embed-MedSAM (Ours)	75.86	146.83	10.6M	33.5G	28.78

The bold values indicate the best results among the compared methods.

In addition, we conducted the Wilcoxon signed-rank test on all datasets to compare Embed-MedSAM with other baseline models. For this purpose, we performed five additional independent experiments for both each baseline model and Embed-MedSAM (the same applies to the following experiments). The results showed statistically significant differences (p < 0.05), as indicated by “^a” in the p-value row of the table, further confirming the superior accuracy and robustness of our proposed method.

Figure 4 shows the visualized segmentation results of several mainstream models on six seen-domain datasets. From the figure, it can be observed that on most benchmark datasets, the segmentation results of ESP-MedSAM and the proposed Embed-MedSAM are similar, and both are close to the Ground Truth. On the ${\mathcal{S}}^6$ and ${\mathcal{S}}^7$ datasets, the segmentation results of Embed-MedSAM are closer to the Ground Truth compared to ESP-MedSAM, showing better boundary capturing and detail preservation. It is also worth noting that Embed-MedSAM uses only one-third of the parameters and half of the inference complexity of ESP-MedSAM, yet achieves comparable or even better segmentation performance, further demonstrating a good balance between segmentation accuracy and computational efficiency.

Fig. 4. Visual comparison of Embed-MedSAM with other fully automated segmentation architectures and lightweight SAM series models on six source domain segmentation tasks (Internal datasets)).

Each row corresponds to a medical imaging modality listed in Table 1, ranging from ${\mathcal{S}}^1$ to ${\mathcal{S}}^9$. The comparison includes traditional segmentation models (such as nnUNet, TransUNet, and MADGNet), lightweight SAM-based models (including Efficient-SAM, SAM-Lightening, and ESP-MedSAM), and our proposed Embed-MedSAM. The rightmost column shows the ground truth annotations. Red boxes highlight segmentation errors relative to the ground truth. As shown, Embed-MedSAM produces segmentation results that are consistently close to the ground truth across all modalities.

Point Prompt lightweight SAM series models. Finally, we further evaluated the lightweight promptable SAM models by introducing point prompts. The results show that, compared to the no-prompt setting, the point-prompt mode-despite being more labor-intensive-does improve segmentation accuracy for these models. However, this improvement comes at the cost of increased professional requirements for end users. Moreover, even with point prompts, none of the lightweight SAM models surpassed the performance of our proposed no-prompt model on any dataset. This demonstrates that our method not only eliminates the need for manual interaction-thereby enhancing usability and accessibility-but also achieves superior segmentation performance. These findings further highlight the practicality and potential of Embed-MedSAM for real-world medical imaging applications. It is also important to note that ESP-MedSAM was not included in this experiment, as it does not support prompt-based inputs.

Diffusion-based medical image segmentation framework architectures. To further validate the overall performance of the proposed lightweight Embed-MedSAM model in medical image segmentation tasks, we include diffusion-based segmentation frameworks-known for their impressive recent performance-as additional baselines. Extensive evaluations were conducted on the ${\mathcal{S}}^4,{\mathcal{S}}^5,{\mathcal{S}}^8$, and ${\mathcal{S}}^{10}$ datasets listed in Table 1, with the results summarized in Table 3.

Table 3.

Comparison of Embed-MedSAM with diffusion-based fully automated segmentation architectures in universal medical image segmentation (source domains, internal datasets)

	${\mathcal{S}}^4$		${\mathcal{S}}^5$		${\mathcal{S}}^8$		${\mathcal{S}}^{10}$			Efficiency
Method	Dice ↑	IoU ↑	Dice ↑	IoU ↑	Dice ↑	IoU ↑	Dice ↑	IoU ↑	HD95 ↓	Params ↓	FLOPs ↓
EnsemDiff59	94.3a	87.8	84.2a	74.4a	87.5a	82.9	88.7	80.9a	10.85	23M	2203G
SegDiff60	92.6a	85.2a	82.5a	71.9a	85.4a	78.8a	85.7a	77.0a	14.31	23M	2399G
MedsegDiff8	95.1a	87.6	85.9a	76.2a	-	-	88.9	81.2a	10.41	25M	1770G
MedsegDiff8+TransUNet53	91.8a	84.5a	82.1a	72.6a	-	-	86.1a	78.0a	13.88	>25M	>1770G
MedSegDiff-V226	96.7	88.9	87.9	80.3	84.4a	77.2a	90.8	83.4	7.53	46M	983G
Embed-MedSAM (Ours)	96.3	87.8	88.4	79.6	88.9	80.4	88.6	82.8	13.96	10.6M	33.5G
(Standard Deviation)	± 0.4	± 0.6	± 0.6	± 0.2	± 0.5	± 0.3	± 0.5	± 0.4	± 1.3	-	-

Bold indicates the best result, and underline indicates the second-best result.

^aScores marked with “^a” indicate that Embed-MedSAM has a statistically significant advantage (p < 0.05) over the specialist models.

As shown in the table, Embed-MedSAM achieves slightly lower segmentation accuracy than MedSegDiff-V2; however, the performance gap is minimal. For instance, on the ${\mathcal{S}}^4$ dataset, its Dice score is only 0.4% lower. Notably, the inference complexity of MedSegDiff-V2 reaches 983 GFLOPs-nearly 30 times higher than that of Embed-MedSAM (33.5 GFLOPs). This extreme computational demand significantly limits the practicality of diffusion models in real-world applications. In contrast, Embed-MedSAM offers a much lower computational cost while maintaining competitive accuracy, thus achieving a superior trade-off between efficiency and performance.

We further evaluated Embed-MedSAM on two large-scale medical image segmentation benchmarks, BTCV, with results shown in Table 4. Consistent with the results in Table 3, Embed-MedSAM achieves segmentation accuracy comparable to MedSegDiff-V2, while its inference complexity is only 1/30 of that of MedSegDiff-V2. This further confirms the practical applicability of Embed-MedSAM for deployment in resource-constrained environments.

Table 4.

Comparison of Embed-MedSAM with diffusion-based fully automated segmentation architectures on the BTCV dataset with multiple anatomical structures (source domains, internal datasets)

	Different Parts
Method	Spleen	R.Kid	L.Kid	Gall.	Eso.	Liver	Stom.	Aorta	IVC	Veins	Panc.	AG	Ave	FLOPs ↓
TransUNet53	95.2a	92.7	92.9a	66.2a	75.7a	96.9	88.9a	92.0	83.3a	79.1a	77.5a	63.7a	83.8a	865.31G
UNetr61	96.8	92.4a	94.1a	75.0a	76.6a	97.1	91.3a	89.0	84.7a	78.8a	76.7a	74.1a	85.6a	-
Swin-UNetr62	97.1	93.6	94.3	79.4a	77.3a	97.5	92.1a	89.2	85.3a	81.2a	79.4a	76.5a	86.9a	-
nnUNet7	94.2a	89.4a	91.1a	70.4a	72.3a	94.8a	82.4a	87.7a	78.2a	72.0a	68.0a	61.6a	80.2a	554.40G
EnsemDiff59	93.8a	93.1	92.4a	77.2a	77.1a	96.7	91.0a	86.9a	85.1a	80.2a	77.1a	74.5a	85.4a	2203G
SegDiff60	95.4a	93.2	92.6a	73.8a	76.3a	95.3a	92.7a	84.6a	83.3a	79.6a	78.2a	72.3a	84.7a	2339G
MedsegDiff8	97.3	93.0	95.5	81.2a	81.5a	97.3	92.4a	90.7	86.8a	82.5a	78.8a	77.9a	87.9	1770G
MedSegDiff8 + TransUNet53	91.2a	87.6a	84.6a	64.5a	71.8a	94.7a	82.4a	87.6a	71.5a	77.5a	67.2a	61.8a	78.5a	>1770G
MedSegDiff-V226	97.8	94.1	96.3	84.8	81.8	98.5	94.0	92.8	86.9	82.3a	83.1	81.7	89.5	983G
Embed-MedSAM (Ours)	96.7	93.2	94.8	84.0	82.6	96.6	93.8	88.8	87.7	83.8	82.3	80.8	88.8	33.5G

Bold indicates the best result, and underline indicates the second-best result.

^aEmbed-MedSAM has a statistically significant advantage (p < 0.05) over the specialist models.

Generalization ability on unseen target domains

To further assess the generalization capability of the proposed model on unseen target domains, we design a cross-dataset generalization evaluation experiment. Specifically, we reuse the models in the experiments from Table 2, ensuring that they have not been exposed to the six target datasets-${\mathcal{T}}^1,{\mathcal{T}}^2,{\mathcal{T}}^3,{\mathcal{T}}^6,{\mathcal{T}}^7$, and ${\mathcal{T}}^9$-at any stage of training. This setup creates a rigorous unseen domain testing scenario (external datasets). For instance, the model trained on the ${\mathcal{S}}^1$ dataset in Table 2 is directly evaluated on the corresponding ${\mathcal{T}}^1$ dataset, allowing us to examine its ability to generalize across different data sources and domain shifts.

The final evaluation results are presented in Table 5. As shown in the table, the findings are consistent with those reported in Table 2. Specifically, Embed-MedSAM outperforms all other no-prompt lightweight SAM models across all unseen datasets, with the exception of ESP-MedSAM. Although ESP-MedSAM achieves slightly higher accuracy on a few datasets, the improvements are marginal. More importantly, it incurs a significantly higher computational cost-its parameter count is approximately three times that of our model, and its FLOPs are nearly twice as high. Furthermore, ESP-MedSAM cannot be deployed on small embedded devices such as the iPhone 14, which severely limits its practical applicability.

Table 5.

Comparison of Embed-MedSAM with other state-of-the-art frameworks in universal medical image segmentation (unseen domains, external datasets)

	Manual	Dermoscopy ${\mathcal{T}}^1$		X-ray ${\mathcal{T}}^2$		Fundus ${\mathcal{T}}^3$		Colonoscopy ${\mathcal{T}}^6$		Ultrasound ${\mathcal{T}}^7$		Microscopy ${\mathcal{T}}^9$		Efficiency
Methods	Prompt	Dice ↑ ↑	HD ↓	Dice ↑ ↑	HD ↓	Dice ↑ ↑	HD ↓	Dice ↑ ↑	HD ↓	Dice ↑ ↑	HD ↓	Dice ↑ ↑	HD ↓	Params ↓	FLPOs ↓
Subtable A. With mainstream fully automatic segmentation architectures
U-Net13	✗	87.00a	140.38a	69.98a	223.70a	61.21a	174.97a	32.60a	461.64a	39.17a	336.98a	46.57a	295.62a	13.40M	497.91G
U-Net++52		87.87a	110.08a	71.29a	222.65a	62.79a	180.01a	36.43a	466.92a	41.30a	343.66a	47.77a	281.42a	9.16M	558.46G
Att-UNet14		88.66a	126.35a	72.41a	234.26a	65.18a	118.10a	35.56a	420.60a	41.89a	357.08a	48.70a	280.56a	8.72M	269.07G
nnUNet7		89.56a	105.53a	74.48a	211.68a	65.00a	107.82a	36.00a	490.92a	43.87a	269.04a	49.19a	283.67a	34.29M	554.40G
H2Former15		87.79a	146.24a	73.34a	219.59a	65.68a	115.71a	34.72a	480.78a	42.46a	275.80a	53.86a	294.98a	33.71M	536.96G
TransUNet53		89.26a	108.48a	80.11a	193.36a	66.23a	112.06a	42.79a	339.14a	44.84a	267.26a	54.22a	282.64a	109.54M	865.31G
ADS18		87.83a	129.46a	76.92a	180.72a	62.45a	172.16a	37.36a	454.70a	43.19a	276.30a	51.06a	281.34a	14.07M	503.59G
CIAug17		88.44a	127.62a	68.22a	257.36a	65.65a	120.87a	39.07a	387.02a	41.50a	267.40a	53.92a	286.23a	14.26M	500.23G
MADGNet16		89.71a	96.86a	84.11a	168.58a	66.88a	119.41a	44.32a	365.84a	44.91a	264.61a	59.29a	278.56a	31.40M	210.65G
Embed-MedSAM (Mean)		91.59	74.34	91.90	94.42	79.43	85.70	66.72	148.59	60.72	226.40	64.77	251.51	10.6M	33.5G
(Standard Deviation)		± 0.38	± 2.02	± 0.29	± 2.23	± 0.72	± 2.45	± 0.50	± 5.57	± 0.56	± 3.97	± 0.66	± 2.40	-	-
Subtable B. With no external prompt lightweight SAM series models
MobileSAM54	✗	85.61a	304.75a	87.60a	277.89a	2.23a	253.99a	20.90a	470.94a	31.03a	342.10a	6.44a	375.52a	9.79M	39.71G
RepViT-SAM55		84.69a	283.65a	86.18a	290.07a	2.30a	283.22a	17.64a	491.23a	26.64a	409.92a	5.96a	383.02a	9.57M	23.64G
EdgeSAM56		84.31a	334.93a	87.07a	277.67a	2.37a	266.70a	17.64a	489.03a	30.01a	382.82a	5.83a	476.44a	9.60M	22.10G
EfficientSAM-S33		88.44a	281.54a	88.24a	281.52a	6.75a	298.82a	46.99a	474.12a	39.44a	468.81a	12.81a	419.71a	25.38M	32.51G
EfficientViT-SAM57		89.29a	177.12a	89.33a	287.08a	9.72a	256.58a	58.77a	357.63a	46.08a	501.00a	14.85a	442.26a	34.80M	89.10G
SAM-Lightening58		89.41a	168.96a	87.62a	300.50a	4.26a	251.03a	46.94a	419.96a	43.68a	446.36a	12.01a	439.31a	19.26M	52.46G
ESP-MedSAM25		91.45	78.13a	92.24	91.30	79.68	83.84	65.96	160.13a	61.62	199.74	64.21	248.47	28.46M	55.88G
Embed-MedSAM (Mean)		91.59	74.34	91.90	94.42	79.43	85.70	66.72	148.59	60.72	226.40	64.77	251.51	10.6M	33.5G
(Standard Deviation)		± 0.38	± 2.02	± 0.29	± 2.23	± 0.72	± 2.45	± 0.50	± 5.57	± 0.56	± 3.97	± 0.66	± 2.40	-	-
Subtable C. With point prompt lightweight SAM series models
MobileSAM54	Point	90.37a	87.28a	89.14a	140.02	54.10a	160.16a	32.77a	399.29a	38.68a	310.52a	16.17a	452.18a	9.79M	39.71G
RepViT-SAM55		90.63a	84.74	88.17a	148.25a	55.72a	133.82a	27.76a	381.39a	33.58a	301.85a	15.20a	438.64a	9.57M	23.64G
EdgeSAM56		90.38a	86.32a	88.01a	147.62a	56.03a	136.48a	28.17a	433.08a	37.51a	291.65a	12.11a	501.14a	9.60M	22.10G
EfficientSAM-S33		90.80a	89.34a	89.00a	150.07a	69.20a	98.05a	56.50a	218.22a	52.95a	233.98a	25.14a	358.13a	25.38M	32.51G
EfficientViT-SAM57		91.14a	85.05a	89.73a	145.84a	72.12a	116.89a	61.67a	179.99a	58.58a	183.63a	34.24a	332.94a	34.80M	89.10G
SAM-Lightening58		90.85a	89.96a	89.63a	144.84a	67.38a	98.85a	48.12a	210.20a	54.75a	244.37a	23.69a	423.15a	19.26M	52.46G
Embed-MedSAM (Mean)	✗	91.59	74.34	91.90	94.42	79.43	85.70	66.72	148.59	60.72	226.40	64.77	251.51	10.6M	33.5G
(Standard Deviation)		± 0.38	± 2.02	± 0.29	± 2.23	± 0.72	± 2.45	± 0.50	± 5.57	± 0.56	± 3.97	± 0.66	± 2.40	-	-

Bold indicates the best result, and underline indicates the second-best result.

^aScores marked with “^a” indicate that Embed-MedSAM has a statistically significant advantage (p < 0.05) over the specialist models.

For promptable SAM models, although the use of point prompts leads to moderate accuracy improvements, none of them surpass the performance of our no-prompt Embed-MedSAM on any unseen dataset. This demonstrates that Embed-MedSAM exhibits stronger generalization and robustness without relying on manual input.

In addition, Fig. 5 shows the segmentation results of several models on various unseen-domain tasks. The visual comparison clearly indicates that most models have limited generalization ability when applied to unseen data, leading to a significant drop in performance. However, ESP-MedSAM and the proposed Embed-MedSAM still show similar results to those in Fig. 4, with segmentation outputs close to the Ground Truth. Considering that Embed-MedSAM uses only half the parameters and one-third of the inference complexity of ESP-MedSAM, yet achieves comparable or even better visual segmentation performance, this further confirms its superior balance between segmentation accuracy and computational efficiency in unseen-domain tasks.

Fig. 5. Visual comparison of Embed-MedSAM with other fully automated segmentation architectures and lightweight SAM series models on six unseen domain segmentation tasks (External datasets)).

Each row corresponds to a medical imaging modality listed in Table 1, ranging from ${\mathcal{T}}^1$ to ${\mathcal{T}}^9$. The comparison includes traditional segmentation models (such as nnUNet, TransUNet, and MADGNet), lightweight SAM-based models (including Efficient-SAM, SAM-Lightening, and ESP-MedSAM), and our proposed Embed-MedSAM. The rightmost column shows the ground truth annotations. Red boxes highlight segmentation errors relative to the ground truth. As shown, Embed-MedSAM produces segmentation results that are consistently close to the ground truth across all modalities.

Ablation study

Key Components. In this section, we design a series of ablation experiments focusing on the three key modules of Embed-MedSAM. The goal is to evaluate the independent contribution of each module to the model’s performance. The final results are presented in Table 6. It should be noted that EKD (Encoder Knowledge Distillation) in the table refers not only to the distillation process itself, but also implies that the image encoder has been replaced by a lightweight structure. Therefore, EKD is marked as enabled in all experiments from Exp 2 to Exp 5.

Table 6.

Ablation study of the three key modules of Embed-MedSAM on general medical image segmentation tasks

	IDE			Seen		Unseen
Experiment	MIP	EKD	SPT	Dice ↑	HD ↓	Dice ↑	HD ↓
Exp 1. MedSAM (Baseline)				85.67	97.56	72.29	179.66
Exp 2.		✓		82.73	113.29	69.22	185.91
Exp 3.	✓	✓		84.36	107.84	71.78	181.02
Exp 4.		✓	✓	88.52	85.30	75.10	149.70
Exp 5. Embed-MedSAM (Ours)	✓	✓	✓	89.64	79.04	75.86	146.82

When SPT is enabled, the model uses a self-prompting mechanism and performs inference without external prompts. When SPT is not enabled, the model is trained normally on the downstream task dataset, and no external prompts are introduced during the inference and evaluation process.

Specifically, compared to the original MedSAM baseline (Exp 1), the complete version of Embed-MedSAM (Exp 5) achieves significant performance improvements on both seen and unseen datasets. On the seen datasets, the average Dice score increases from 85.67% to 89.64%, and the Hausdorff Distance (HD) decreases from 97.56 to 79.04. On the unseen datasets, the Dice score improves from 72.29% to 75.86%, and HD drops from 179.66 to 146.82. Despite these improvements, Embed-MedSAM has far fewer parameters and lower FLOPs compared to MedSAM, showing strong advantages in lightweight design.

In the encoder distillation stage, the Masked Image Pretraining (MIP) module plays a positive role in enhancing the model’s representation ability and accuracy. For example, in Exp 3, after introducing MIP, the Dice score on the seen datasets increases from 82.73% (Exp 2) to 84.36%, with a clear reduction in HD. This shows the importance of MIP during the early stage of feature learning.

In this experiment, the MedSAM baseline does not use any external prompts (such as points or boxes) during inference, which limits its full performance. Once the Self Prompt Training (SPT) module is enabled, the model no longer relies on manual prompts. Instead, it generates pseudo-prompts from the initial prediction and uses them to guide a more accurate second prediction. The results show that in all configurations where SPT is enabled, the Dice score improves significantly. This confirms that the self-prompting mechanism greatly enhances the model’s adaptability and usability under prompt-free conditions.

Epochs of Masked Image Pretraining. In the default setting, we set the number of training epochs for Masked Image Pretraining (MIP) in the Image Encoder Distillation (IED) stage to no less than 10. To validate this choice, we conducted an ablation study. As shown in Fig. 6, as the number of training epochs increases from 1 to 10, the segmentation accuracy of Embed-MedSAM on both seen and unseen domains first improves significantly and then gradually stabilizes. This indicates that moderate pretraining on general-domain images may be helps the model learn structural representations with good generalization and transferability.

Fig. 6. Effect of MIP epochs on segmentation performance in seen and unseen domains.

The x-axis indicates the number of Masked Image Pretraining (MIP) epochs during the IDE stage. The left y-axis corresponds to Seen Dice (blue triangles), and the right y-axis corresponds to Unseen Dice (red squares).

However, when the training epochs further increase to between 11 and 15, model performance begins to plateau. In some tasks, a slight decline is observed. We speculate that excessive training may cause the lightweight encoder to overfit the visual patterns in the general domain. This could reduce its ability to capture structural features in medical images, thereby affecting generalization on downstream segmentation tasks.

In the remaining training stages of this study, we primarily adopt an early stopping strategy to dynamically control the training process. The actual number of training epochs varies across different tasks. Therefore, it is unnecessary to perform additional ablation studies on fixed epoch settings for these stages.

Loss Weight Sensitivity Analysis. In this study, we design composite loss functions for both the EKD (Encoder Knowledge Distillation) and SPT (Self Prompt Training). Each loss function integrates multiple supervision terms and introduces corresponding loss weighting factors (λ₁ to λ₄)to balance the contributions of different sub-tasks. To evaluate the impact of various weight configurations on model performance, we conduct ablation studies and present the results in Table 7.

Table 7.

Ablation study on the weights of individual loss components in the composite loss function

(a) Weight λ1.			(b) Weight λ2.			(c) Weight λ3.			(d) Weight λ4.
	Seen	Unseen		Seen	Unseen		Seen	Unseen		Seen	Unseen
λ1	Dice ↑	Dice ↑	λ2	Dice ↑	Dice ↑	λ3	Dice ↑	Dice ↑	λ4	Dice ↑	Dice ↑
0	76.29	61.38	0	89.03	75.38	0	-	-	0	-	-
1 × 10−5	80.29	67.20	0.1	89.64	75.86	0.4	86.46	71.79	0.1	86.29	73.77
3 × 10−5	85.32	65.92	0.2	87.69	72.24	0.8	89.64	75.86	0.2	89.64	75.86
5 × 10−5	89.64	75.86	0.3	84.39	73.40	1.2	88.43	73.29	0.3	86.71	73.92
7 × 10−5	88.72	70.83	0.4	82.81	72.77	1.6	85.37	73.38	0.4	84.37	69.51

The Dice score reported here is the average result evaluated across multiple datasets.

The bold values indicate the best results among the compared methods.

As shown in Table 7, different weight settings for each loss component have a significant impact on the Dice scores in both seen and unseen domains. For the semantic consistency loss (λ₁), the best performance is achieved when the weight is set to 5 × 10⁻⁵, with Dice scores of 89.64 (Seen) and 75.86 (Unseen). This shows that a moderate KL divergence constraint helps stabilize semantic alignment in self-prompt training and improves generalization. For the geometric alignment loss (λ₂), the optimal setting is λ₂ = 0.1, confirming that supervising the alignment between the pseudo mask and the ground truth box improves segmentation quality. In terms of segmentation accuracy, the best performance is achieved when the weight of the Dice loss (λ₃) is set to 0.8 and the weight of the BCE loss (λ₄) is set to 0.2. When these weights are too small or too large, the model performance decreases, showing that proper weight selection is important for stable optimization.

Overall, the default setting provides a good balance among the four loss terms. Each term contributes supervision from a different aspect: semantic consistency (${\mathcal{L}}_{\mathrm{KL}}$), geometric alignment (${\mathcal{L}}_{\mathrm{Box}}$), and pixel-level accuracy (${\mathcal{L}}_{\mathrm{Dice}}$ and ${\mathcal{L}}_{\mathrm{BCE}}$). This multi-level supervision helps the model achieve good generalization and segmentation performance across diverse medical image datasets.

It should be clarified that the default weight settings mentioned above represent the overall best-performing configuration across all datasets. However, in practical applications, these weights should be dynamically adjusted based on specific training feedback to further optimize the model’s performance on the target downstream task.

Embedded deployment

In the previous experiments, the training and evaluation of our model were mainly conducted on a high-performance NVIDIA 2080Ti GPU platform. However, since the proposed Embed-MedSAM is designed for deployment on resource-constrained edge devices, its feasibility on embedded platforms is particularly important. Therefore, in this section, we conduct an empirical evaluation of the model’s inference efficiency on a low-computation mobile platform-iPhone 14. Specifically, we followed the deployment method described in the EdgeSAM paper. The model was deployed on the iPhone 14 using Core ML Tools as the compiler. The model throughput was measured using the benchmarking tool provided by Xcode. The final results are shown in Table 8.

It is important to note that deployment details and hardware differences can lead to significant variation in results. To ensure fair comparison, we used the FPS values reported in the original papers for all baseline models, instead of values obtained from local reproduction.

As shown in the table, the original SAM, MedSAM, and the less efficient variants EfficientSAM-S and ESP-MedSAM all encountered out-of-memory (OOM) issues on the iPhone 14 (local independently). This indicates that these models are not suitable for deployment on memory-limited mobile or embedded devices due to their large parameter size and high computational demand. Among the lightweight SAM variants that can run successfully, we observed that their segmentation performance on unseen domains is clearly lower than that of the proposed Embed-MedSAM. Although models such as MobileSAM and RepViT-SAM meet the constraints of low resource usage, they exhibit poor generalization ability and low Dice scores, making them less reliable in real-world medical segmentation tasks.

In contrast, Embed-MedSAM can be efficiently deployed on devices with limited resources such as the iPhone 14, while still maintaining high segmentation accuracy on unseen domain tasks. This demonstrates that our model achieves a good balance between lightweight design and robust performance, making it highly suitable for real-world medical image analysis on edge devices.

In addition, we report the runtime resource usage of Embed-MedSAM on the iPhone 14 during one hour of manually operate and perform continuous inference. The model consumed approximately 25% of battery power, reached a peak memory usage of 1.9GB (exceeded), and maintained a CPU usage of around 30% (exceeded). The surface temperature of the device remained stable, with only slight warming observed. These results further demonstrate that Embed-MedSAM can run stably with low resource consumption on small embedded devices, confirming its potential for real-world edge deployment.

Discussion

This paper presents Embed-MedSAM, a lightweight and prompt-free segmentation model tailored for medical imaging applications in real-world, resource-constrained environments-such as conflict zones, rural areas, and underdeveloped regions.

Embed-MedSAM addresses two major limitations of existing SAM-based models. First, many current models require substantial memory and computational resources. To overcome this, we replace the original MedSAM encoder with a lightweight RepViT backbone and adopt a two-stage distillation strategy. The encoder is first pretrained on ImageNet-21K using masked image reconstruction, then fine-tuned on the COSMOS 1050K medical dataset using MSE loss. This enables the lightweight encoder to retain the semantic representation capabilities of MedSAM while significantly reducing model size and inference latency.

Secondly, we completely eliminate the reliance on manual prompts during inference by introducing a self-prompting mechanism. Without any prompt input, the model uses the segmentation mask generated from the first forward pass to derive a pseudo box prompt via its minimum bounding rectangle. This pseudo prompt is then refined by the prompt encoder to guide a second, more accurate prediction. To suppress error accumulation from inaccurate pseudo-prompts, we incorporate a KL divergence loss during training to enforce semantic consistency between the two predictions. Additionally, an L1 loss is applied between the pseudo box and the ground truth mask box to enhance geometric alignment. At the mask level, a combination of Dice loss and Binary Cross-Entropy (BCE) loss is used to optimize both the overlap and pixel-wise accuracy between the predicted and ground truth masks. This framework eliminates manual prompt dependency in both training and inference stages, ensuring strong generalizability and deployment efficiency.

Experimental results across 17 datasets spanning over 7 medical imaging modalities demonstrate that Embed-MedSAM achieves strong performance in both source (seen) and target (unseen) domains. The model offers a favorable trade-off between segmentation accuracy and deployment efficiency and is capable of running at nearly 30 FPS on devices such as the iPhone 14. These results suggest that universal medical image segmentation is feasible without reliance on high-performance hardware and manual inputs.

Although Embed-MedSAM is lightweight and efficient during inference, its initial training still requires a certain level of computational resources. Specifically, the first-stage encoder distillation adopts a two-step strategy and requires more than one day of training on four NVIDIA A800 GPUs. In contrast, the second-stage Self Prompt Guided Mask Decoder Training focuses on adapting the model to specific downstream tasks, with significantly lower training costs. For example, on the REFUGE dataset, this stage takes only about 30 minutes using two NVIDIA 2080Ti 11GB GPUs. Therefore, after completing the first-stage distillation training once, users can easily fine-tune the model for their own medical image segmentation tasks using low-cost hardware. This paradigm greatly enhances the practicality and scalability of Embed-MedSAM in real-world medical applications.

In conclusion, Embed-MedSAM provides an efficient and accessible foundation model for medical image segmentation in low-resource environments. Its lightweight design, prompt-free inference, and robust generalization make it well-suited for deployment in real-world clinical scenarios with limited infrastructure and expertise. With a one-time training investment, the model supports broad deployment and has the potential to improve diagnostic accessibility and efficiency in underserved regions.

Methods

Preliminaries of MedSAM

SAM is a highly versatile framework for image segmentation, capable of handling diverse tasks with minimal human input. Based on SAM, MedSAM²¹ extends its capabilities to the medical domain by training on a large-scale medical imaging dataset. The MedSAM architecture includes three main components: an image encoder, a prompt encoder, and a mask decoder (see Fig. 1).

Image Encoder. The image encoder is the most computationally intensive component of MedSAM. It adopts a hierarchical architecture based on Vision Transformer (ViTDet)²⁷, which enables the model to capture multi-scale features. The encoder takes a fixed-size input of (3, 1024, 1024) and outputs a feature map of size (256, 64, 64). This powerful backbone allows MedSAM to extract deep visual features effectively, supporting robust performance across diverse image domains.

Prompt Encoder. MedSAM supports several types of prompts, including points, boxes, masks, and text, to guide accurate mask generation. The prompt encoder processes these inputs to enhance segmentation. Sparse prompts, such as points and boxes, are encoded using random positional embeddings along with a type indicator token to distinguish between prompt types. For more complex prompts-like masks and free-form text-MedSAM uses a small CNN and the CLIP text encoder²⁸, respectively. In this study, we focus on point and box prompts, which are commonly used in interactive segmentation tasks.

Mask Decoder. The mask decoder generates the final segmentation output. It takes as input the image features from the encoder, embedded prompts (points or boxes), mask prediction tokens, and an IoU (Intersection over Union) prediction token. These inputs are processed through a lightweight two-way transformer. Each mask token is converted into a dynamic linear classifier that computes foreground probabilities at each spatial location. The IoU token provides a confidence score for the predicted mask. This architecture enables MedSAM to generate accurate and adaptable masks, even under challenging conditions.

Large-Scale Medical Image Segmentation Dataset. A major factor behind MedSAM’s strong performance is its training on a large-scale medical image dataset. This dataset includes 1,570,263 image-mask pairs, spanning 10 imaging modalities and over 30 cancer types. As the largest segmentation dataset to date, it significantly improves MedSAM’s zero-shot generalization across diverse medical tasks.

Research objective

The primary objective of this study is to develop a medical image segmentation model based on the advanced MedSAM framework, with two key goals: enabling deployment on resource-constrained embedded devices and eliminating the need for prompt inputs during inference. The model must significantly reduce computational cost during inference while preserving the high segmentation accuracy achieved by the original MedSAM.

It is important to note that the image encoder in MedSAM contains over 600 million parameters, whereas the combined parameter count of the prompt encoder and mask decoder is under 4 million. Therefore, this study addresses two major challenges: (1) how to replace the large-scale image encoder with a lightweight alternative that maintains MedSAM-level segmentation accuracy, enabling efficient deployment on edge devices; and (2) how to remove the dependency on prompt inputs at inference time.

Lightweight MedSAM

We introduced a lightweight image encoder based on RepViT²⁴. RepViT is a convolutional neural network (CNN) derived from MobileNetV3²⁹, and incorporates design efficiencies from modern lightweight Vision Transformers (ViTs). It achieves superior performance compared to other lightweight ViT models of similar size and demonstrates lower latency on mobile devices such as the iPhone 12.

In this work, we selected the RepViT-M1.0 variant, which offers significantly faster CPU inference speed than the baseline model using the TinyViT backbone³⁰. The macro structure of our encoder is illustrated in Fig. 3, consisting of a stem, two downsampling layers, and an encoding layer. Between these stages, we insert multiple RepViT blocks in a 1:1:8 ratio. RepViT also adopts a structural re-parameterization technique³¹, which reduces both computational and memory overhead, enabling faster inference on resource-limited devices without sacrificing accuracy.

To ensure that the lightweight MedSAM achieves performance comparable to the original MedSAM, while removing its reliance on prompt inputs, we adopt a two-stage training strategy inspired by the “divide-and-conquer training paradigm” proposed by Zhang et al.³². Specifically, we decompose the training process into two independent sub-tasks: image encoder distillation and prompt guided mask decoder training.

Image encoder distillation

After replacing the original large-scale image encoder with a lightweight counterpart, conventional methods often perform knowledge distillation directly on the target domain (e.g., medical imaging data) to transfer the original encoder’s feature extraction and semantic understanding capabilities. However, this approach can easily lead to overfitting due to the limited capacity of lightweight models.

To mitigate this issue, we propose a novel two-step distillation strategy, as shown in Algorithm 1. In the first step, we perform masked image pretraining on the part of ImageNet-21K dataset [33], using reconstruction supervision from MedSAM to transfer general visual knowledge. In the second step, we apply optimization of Mean Squared Error (MSE) loss [23] on the subset of large-scale medical dataset COSMOS 1050K. This guides the lightweight encoder to align with MedSAM’s semantic features and effectively inherit its domain adaptability in medical image segmentation tasks.

Masked image pretraining

In this fine-tuning stage, inspired by MAE and EffficientSAM³³, we introduce a MedSAM-guided masked image pretraining strategy. As shown in Fig. 2a step 1, the lightweight encoder extracts latent features from unmasked tokens, while the cross-attention decoder reconstructs the masked tokens using both the unmasked features and the target features extracted by MedSAM. This enables the encoder to learn MedSAM-style representations on ImageNet-21K without requiring any medical annotations.

Algorithm 1

Procedure of two-step distillation strategy

1: Input: General ImageNet-21K dataset ${\mathcal{D}}_{\mathrm{gen}}$, COSMOS 1050K medical dataset ${\mathcal{D}}_{\mathrm{med}}$ with pseudo masks; pretrained MedSAM encoder f^sam

2: Output: Trained lightweight encoder g^e and adapted decoder $\mathcal{D}$

3: Step 1: Masked image pretraining

4: for each image $\mathbf{x}\in {\mathcal{D}}_{\mathrm{gen}}$ do

5:  Apply random token masking $\to \mathcal{M},\mathcal{U}$

6:  Encode unmasked tokens: ${\mathbf{z}}_{\mathcal{U}}=g^e\left({\left\{{\mathbf{x}}_i\right\}}_{i\in \mathcal{U}}\right)$

7:  Decode masked tokens: ${\mathbf{z}}_{\mathcal{M}}=g^d\left({\left\{{\mathbf{x}}_i\right\}}_{i\in \mathcal{M}}\right)$

8:  Merge: $\mathbf{z}=\varphi \left({\mathbf{z}}_{\mathcal{U}}\oplus {\mathbf{z}}_{\mathcal{M}}\right)$

9:  Project: $\widehat{f}=h^{\theta }\left(\mathbf{z}\right)$

10:   Compute MedSAM target: f^sam(x)

11:   Compute reconstruction loss: $L=\frac{1}{N}{\sum }_{j=1}^N\parallel {\widehat{f}}_j-f_j^{\mathrm{sam}}{\parallel }^2$

12:   Update g^e, g^d, and h^θ to minimize L

13: end for

14: Step 2: Encoder knowledge distillation

15: for each image $\mathbf{x}\in {\mathcal{D}}_{\mathrm{med}}$ do

16:   Extract features: f^T(x) from MedSAM, f^S(x) from g^e

17:   Compute MSE loss: $L_{\mathrm{Distill}}=\frac{1}{N}{\sum }_i\parallel f^T\left({\mathbf{x}}_i\right)-f^S\left({\mathbf{x}}_i\right){\parallel }^2$

18:   Update g^e to minimize L_Distill

19: end for

Cross-Attention Decoder. Only masked tokens are reconstructed, with the encoder outputs for unmasked tokens serving as reference keys and values. Masked tokens are treated as queries. The outputs of masked and unmasked tokens are then merged and reordered to produce the final latent embedding.

Linear Projection Head. To match the dimensionality of MedSAM features, we apply a lightweight linear projection head to the decoder output, enabling feature-level alignment.

Reconstruction Loss. During each training step, the process involves extracting features through a forward pass in the MedSAM image encoder, followed by forward and backward passes in the MAE model. The reconstruction loss is then computed by comparing the outputs from the MedSAM encoder and the linear projection head in MAE.

Let the MedSAM image encoder be represented as f^sam, and the MAE’s encoder and decoder as g^e and g^d, with weights W_e and W_d respectively. We define the linear projection head as h^θ, with weights W_θ. Suppose the input tokens are represented by ${\left\{{\mathbf{x}}_i\right\}}_{i=1}^N$, where N is the total number of tokens. These tokens are split into unmasked tokens, ${\left\{{\mathbf{x}}_i\right\}}_{i\in \mathcal{U}}$, and masked tokens, ${\left\{{\mathbf{x}}_i\right\}}_{i\in \mathcal{M}}$, based on a predefined masking ratio. We denote the reordering function as ϕ and the merging function as ⊕ .

The target feature output from the MedSAM encoder can be expressed as $f^{\mathrm{sam}}\left(\mathbf{x}\right)=f^{\mathrm{sam}}\left({\left\{{\mathbf{x}}_i\right\}}_{i=1}^N\right)$. The MAE encoder output for unmasked tokens is $g^e\left({\left\{{\mathbf{x}}_i\right\}}_{i\in \mathcal{U}}\right)$, while the decoder generates $g^d\left({\left\{{\mathbf{x}}_i\right\}}_{i\in \mathcal{M}}\right)$. The output from the linear projection head is then:

$$f^h\left(\mathbf{x}\right)=h^{\theta }\left(\varphi \left(g^e{\left\{{\mathbf{x}}_i\right\}}_{i\in \mathcal{U}}\oplus g^d{\left\{{\mathbf{x}}_i\right\}}_{i\in \mathcal{M}}\right)\right).$$ 1

Therefore, our target reconstruction loss can be formulated as,

$$L_{W_e,W_d,W_{\theta }}=\frac{1}{N}\cdot \sum _{j=1}^N\mid \mid f^{\mathrm{sam}}\left(\mathbf{x}\right)-f^h\left(\mathbf{x}\right)\mid {\mid }^2,$$ 2

where N is the number of input tokens, ∣∣ ⋅ ∣∣ denotes a norm. We use l₂ norm for reconstruction loss in our experiments. By minimizing the reconstruction loss, $L_{W_e,W_d,W_{\theta }}$, our encoder g^e is optimized to serve as an image backbone to extract features as MedSAM image encoder. Our encoder, decoder, and linear projection head are optimized to learn context modeling ability from MedSAM image encoder. Optimizing the reconstruction loss on all tokens transfer the knowledge embedded in MedSAM.

Encoder knowledge distillation

In this stage, we use the subset of large-scale medical imaging dataset COSMOS 1050K to perform semantic distillation guided by the MedSAM image encoder. This helps the lightweight model inherit MedSAM’s semantic representation and structural perception abilities in medical images, as shown in Fig. 2a, step 2. The goal is to enable effective knowledge transfer and domain adaptation in the medical field.

Specifically, we apply a Mean Squared Error (MSE) loss to align the feature maps of the teacher and student models at the pixel level. This loss helps the student model learn the intermediate representations of the teacher, which improves its semantic understanding.

The distillation loss is defined as:

$$L_{\mathrm{Distill}}=\frac{1}{N}\sum _{i=1}^N{∥f^T\left({\mathbf{x}}_i\right)-f^S\left({\mathbf{x}}_i\right)∥}^2,$$ 3

where f^T(x_i) and f^S(x_i) denote the feature maps from the MedSAM image encoder (teacher) and the lightweight RepViT image encoder (student), respectively. N is the total number of pixels in the feature map. This loss ensures that the student model effectively mimics the teacher’s feature representation without adding complexity, making it suitable for deployment in resource-constrained environments.

Self Prompt Training Framework

To enhance the decoder’s compatibility with the previously distilled lightweight image encoder and eliminate reliance on manual prompts, we propose a self-prompt-guided decoder training framework on downstream medical image datasets, which inspired by [22].

Conventional models such as SAM and MedSAM rely on externally provided prompts (e.g., boxes, points, or masks) during both training and inference. To remove this dependency, we introduce a plug-and-play self-prompting module, as shown in Fig. 2b.

Let x denote the input image, and let g^e(x) be the image features extracted by the Image encoder. During the initial training step, the decoder generates a coarse segmentation mask ${\widehat{\mathbf{y}}}^{\left(0\right)}$ using a default prompt embedding p⁽⁰⁾ initialized internally:

$${\widehat{\mathbf{y}}}^{\left(0\right)}=\mathcal{D}\left(g^e\left(\mathbf{x}\right),{\mathbf{p}}^{\left(0\right)}\right),$$ 4

Next, we extract the boundary of ${\widehat{\mathbf{y}}}^{\left(0\right)}$ to generate a self-derived box prompt b⁽¹⁾, which is encoded by the prompt encoder g^p into an updated prompt embedding:

$${\mathbf{p}}^{\left(1\right)}=g^p\left({\mathbf{b}}^{\left(1\right)}\right),$$ 5

and a refined prediction is then produced:

$${\widehat{\mathbf{y}}}^{\left(1\right)}=\mathcal{D}\left(g^e\left(\mathbf{x}\right),{\mathbf{p}}^{\left(1\right)}\right).$$ 6

To ensure robust self-prompt learning and accurate segmentation without manual inputs, we adopt a multi-term loss function that supervises different aspects of the training process:

Semantic Consistency. To prevent error accumulation from the self-generated prompt, we introduce a Kullback-Leibler (KL) divergence loss that enforces semantic alignment between the initial coarse prediction and the refined prediction:

$$L_{\mathrm{KL}}=D_{\mathrm{KL}}\left({\widehat{\mathbf{y}}}^{\left(0\right)}\mid \mid {\widehat{\mathbf{y}}}^{\left(1\right)}\right)=\sum _i{\widehat{\mathbf{y}}}_i^{\left(0\right)}\log \left(\frac{{\widehat{\mathbf{y}}}_i^{\left(0\right)}}{{\widehat{\mathbf{y}}}_i^{\left(1\right)}+ϵ}\right)$$ 7

Geometric Alignment. To improve the quality of the pseudo box prompt, we use an L1 loss between the bounding box of the pseudo mask and that of the ground truth mask:

$$L_{\mathrm{Box}}={∥\mathrm{Box}\left({\widehat{\mathbf{y}}}^{\left(0\right)}\right)-\mathrm{Box}\left(\mathbf{y}\right)∥}_1$$ 8

Segmentation Accuracy. To guide the final prediction toward accurate region delineation, we jointly apply Dice loss and Binary Cross Entropy (BCE) loss between the refined mask prediction and the ground truth mask:

$$L_{\mathrm{Dice}}=1-\frac{2{\sum }_i{\widehat{\mathbf{y}}}_i^{\left(1\right)}\cdot {\mathbf{y}}_i}{{\sum }_i{\widehat{\mathbf{y}}}_i^{\left(1\right)}+{\sum }_i{\mathbf{y}}_i+ϵ}$$ 9

$$L_{\mathrm{BCE}}=-\frac{1}{N}\sum _i\left[{\mathbf{y}}_i\log {\widehat{\mathbf{y}}}_i^{\left(1\right)}+\left(1-{\mathbf{y}}_i\right)\log \left(1-{\widehat{\mathbf{y}}}_i^{\left(1\right)}\right)\right]$$ 10

The total training loss is formulated as a weighted sum of all components:

$$L_{\mathrm{total}}={\lambda }_1{L}_{\mathrm{KL}}+{\lambda }_2{L}_{\mathrm{Box}}+{\lambda }_3{L}_{\mathrm{Dice}}+{\lambda }_4{L}_{\mathrm{BCE}}$$ 11

where λ₃ to λ₆ are loss weights empirically selected to balance consistency, alignment, and accuracy. By default, the loss weight parameters are set to: λ₁ = 5 × 10⁻⁵, λ₂ = 0.1, λ₃ = 1, and λ₄ = 1. However, these hyperparameters can be adjusted according to the characteristics of specific downstream medical imaging datasets to achieve better training performance.

Author contributions

Y.Z.: conceptualization, methodology, software, validation, formal analysis, resources, investigation, data curation, writing - original draft, visualization. F.Y.: conceptualization, methodology, software, validation, formal analysis, resources, investigation, data curation, writing- original draft, visualization. X.Y.: conceptualization, methodology, software, validation, formal analysis, resources, investigation, data curation, writing - original draft, visualization. X.L.: validation, resources, writing– review & editing, supervision, project administration, funding acquisition. T.J.: writing – review & editing, supervision. L.Y.Y.: writing – review & editing, supervision. L.Y.: validation, resources, writing – review & editing, supervision, project administration, funding acquisition. All authors have read and approved the manuscript.

Data availability

The ImageNet-21K dataset can be downloaded from https://huggingface.co/datasets/gmongaras/Imagenet21K. The COSMOS 1050K dataset can be accessed via the download link provided at https://github.com/yuhoo0302/Segment-Anything-Model-for-Medical-Images.

Code availability

The Step 1 of the Image Encoder Distillation code is based on the description in EfficientSAM. Step 2 of the Image Encoder Distillation code is implementation of Rep-MedSAM. The Self Prompt Training Framework code refers to SAM-SP, and a simple example code is provided on https://github.com/XinLianSJ/Embed-MedSAM.

Competing interests

The authors declare no competing interests.