- 期刊:npj Digital Medicine
- 发表:2025年9月24日
- 亮点:提出轻量级、无人工提示的 Embed-MedSAM,可在 iPhone 14 上接近 30 FPS。
- 建议重点看:知识蒸馏、自提示训练、精度与计算成本的权衡
Embedded framework for clinical medical image segment anything in resource limited healthcare regions
面向资源有限医疗区域的临床医学图像任意分割嵌入式框架
摘要
便携式成像设备的出现提升了资源有限地区的医学图像采集效率,但医务人员短缺仍制约着及时诊断。为此,我们提出 Embed-MedSAM,一种部署成本低的全自动分割模型。该模型以 MedSAM 为基础,集成了轻量级 RepViT 编码器以降低计算量,并在超过一百万张多模态医学图像上采用两阶段蒸馏,以保留原始模型的视觉表示能力。同时引入自提示机制,由模型生成伪标签来引导细粒度分割。训练过程联合优化 KL 散度与分割损失,以提升无提示条件下的分割精度。在涵盖 7 种成像模态的 17 个基准数据集上,Embed-MedSAM 展现出卓越性能。在无外部提示的情况下,其平均 Dice 分数比排名第二的模型高出近 16%。该模型在 iPhone 14 上运行速度接近 30 FPS,显示出较强的实际部署潜力。
理解1
该模型以 MedSAM 为基础,集成了轻量级 RepViT 编码器以降低计算量,并在超过一百万张多模态医学图像上采用两阶段蒸馏,以保留原始模型的视觉表示能力。
MedSAM 是一个基于"分割一切模型"(SAM)微调出来的医学图像分割模型,效果很好,但模型较大、计算量高,不太适合在手机或便携设备上跑。
编码器:用来从图像中提取特征
换了小编码器后,如果直接训练,模型的能力可能远不如原来那个大模型。为了 "用小模型复现大模型的本事",他们用了"知识蒸馏"。
这里的"蒸馏",可以理解为让大模型当老师,小模型当学生,学生模仿老师对图像的理解。
理解2
同时引入自提示机制,由模型生成伪标签来引导细粒度分割。
原始的 SAM(Segment Anything Model)和 MedSAM 虽然能分割万物,但需要你给它一个提示,比如点一下病灶位置、画个框,它才知道要分割哪里。
但在资源有限地区,希望模型全自动,不能指望医生一个一个去点。这就是"无外部提示"的条件。
"自提示"(self-prompting)就是模型自己给自己生成提示,不再依赖人工点选。
具体做法:模型先对图像做一个初步的、粗略的分割,得到一个大致的区域位置。这个粗略分割就充当了"提示"------相当于模型自己画了个框或者指了个位置。
有了伪标签提供的粗定位后,模型再进行第二次更精细的分割,这就是细粒度分割。
理解3
训练过程联合优化 KL 散度与分割损失,以提升无提示条件下的分割精度。
KL 散度是一个衡量两个概率分布之间差异 的指标。在这里,它用来计算学生模型输出的概率分布 和教师模型输出的概率分布有多大的差距。
训练时,把 KL 散度作为一项损失,目标就是让它越小越好。越小,说明学生模仿老师模仿得越像,大模型的"视觉经验"就被更完整地迁移到了小模型里。
在无提示条件下,模型需要自己生成伪标签来引导分割。联合优化能让伪标签更靠谱(受蒸馏的强视觉表示支撑),同时让细粒度分割更精确(受分割损失直接监督),最终大幅提升全自动分割的精度。
引言
在冲突地区、偏远地区和经济欠发达区域等资源受限环境中,医疗资源的匮乏仍是改善医疗服务的主要障碍。近年来,便携式临床成像设备的研发与部署,大幅提高了这些地区的医疗可及性,使本地机构能更便捷地采集高质量医学图像。
然而,仅靠成像设备并不能解决核心问题。这些地区通常严重缺乏训练有素的专业人员,尤其是擅长解读复杂医学图像的专家。在此类环境中,快速分析图像并检测病灶的能力至关重要。诊断延误可能导致错失治疗时机。因此,在硬件条件受限的情况下,对不同模态的医学图像进行高效分割和诊断,仍是一项迫切且尚未解决的挑战。
为应对这一需求,基于计算机视觉的自动分割技术已被引入医学图像分析中。这些方法旨在识别并勾画出器官、组织等感兴趣区域(ROIs)。近来,自动分割方法因能提升一致性和精度而受到关注。与此同时,深度学习的进步使神经网络在医学图像分割中得到广泛应用,包括传统的卷积神经网络(CNNs)以及近年来的视觉Transformer(ViTs)。
基于CNN的U形架构对低层与高层语义特征之间的关系进行建模,以预测掩码。这类方法在已知领域中能产生准确的分割结果,但在未见领域中的性能往往下降。为解决领域泛化(DG)的挑战,现有方法采用多任务学习、数据增强和领域合成等策略,旨在增加所学特征表示的多样性。然而,由于模型容量有限,许多此类任务特定模型在应用于不同成像模态时,必须从头重新训练。这为构建在多种医学领域具有强泛化能力和通用性的分割模型带来了主要障碍。
视觉Transformer(ViTs)通过自注意力捕获长距离依赖关系,从而解决了这一局限。在大规模数据集上训练时,ViTs得益于更大的模型容量和更少的归纳偏置。近期,Segment Anything Model(SAM)在图像分割方面取得了重大突破。SAM在自然图像中展现出强大的泛化能力,并可通过基于提示的交互在各种场景中表现出色。利用其强大的预训练编码器,多项研究已探索了SAM在通用医学图像分割中的潜力。
然而,尽管MedSAM(如图1所示)是当前SAM系列中专门为医学图像分割设计的模型,但在实际应用中仍面临两大挑战,有待解决。
(1)MedSAM及其变体参数规模极其庞大,通常超过6亿,导致推理过程中计算成本高昂。针对这一问题,现有两种常见方案:(i)为算法配备专用的高性能计算设备;或(ii)将模型部署在云端并远程传输推理结果。然而,在硬件和网络资源有限的地区(如冲突区域、偏远山区或高海拔地区),这两种方案均不切实际,且后者还引发数据隐私和安全方面的顾虑。
(2)MedSAM的推理过程严重依赖用户提供的提示(如点或边界框)来引导分割掩码的生成。尽管基于边界框的提示能产生相对准确的结果,但用户通常需要大致指出图像中病灶或目标结构的位置。这对用户提出了一定程度的医学专业知识要求,限制了模型在不发达或非专业环境中的适用性。
图一
图 1 | MedSAM 框架概述。输入图像首先经过掩码编码器处理以提取图像特征;这些特征随后与提示信息(如掩码、点、矩形框或文本)结合;提示编码器对这些输入数据进行处理;最后,掩码解码器利用图像嵌入向量和提示嵌入向量共同生成最终的分割结果。

为应对上述挑战,我们提出了Embed-MedSAM(嵌入式医学图像分割一切模型),如图2所示。具体而言,Embed-MedSAM是MedSAM²¹的优化扩展,通过两阶段过程实现改进:
(1)我们将MedSAM中计算成本高昂的图像编码器替换为更高效的轻量级RepViT模型²⁴(如图3所示),以显著降低计算成本和推理延迟。为保留MedSAM强大的视觉表征能力,我们设计了一种两步蒸馏策略:首先,在通用领域数据集ImageNet-21K上进行掩码图像预训练,并结合原始MedSAM的重建监督,以迁移通用视觉知识;其次,在大规模医学影像数据集COSMOS 1050K上引入均方误差(MSE)损失的优化,引导轻量级编码器与MedSAM的语义特征对齐,从而有效继承其在医学图像任务中的领域适应性。
(2)我们提出了一种无需人工提示的自提示训练框架用于医学图像分割。该方法利用模型的首次前向预测生成伪掩码,再通过最小外接矩形将其转换为边界框。该伪边界框随后由提示编码器精炼,以引导第二次更准确的预测。为减少伪提示带来的累积误差,我们引入KL散度损失来强制两次预测之间的一致性。此外,在伪边界框与真实边界框之间施加L1损失以改善几何对齐。在掩码层面,使用Dice损失和二元交叉熵损失优化重叠度和像素级精度。该框架在训练和推理过程中均无需人工提供提示,从而能够在真实场景中实现更好的泛化和部署。
图二
图2 | Embed-MedSAM模型的整体技术架构。a 图像编码器蒸馏阶段:a部分展示了用轻量级RepViT编码器替代原始MedSAM编码器的过程。步骤1展示了在ImageNet-21K数据集上进行掩膜图像重建预训练,利用重建损失函数传递通用视觉知识;步骤2则展示了在cosmos 1050K医学数据集上的知识蒸馏过程,通过 MSE 损失函数引导轻量级编码器与MedSAM的语义特征保持一致。b 自主提示训练框架:该阶段旨在消除对人工提示的需求------模型根据初始预测结果生成自指导提示,使用提示编码器进行编码后输入冻结解码器。为抑制伪提示误差的累积,在训练过程中引入KL散度损失以确保两次预测结果间的语义一致性;同时采用伪提示框与真实掩膜框之间的L1损失来提升几何对齐精度。在掩膜层面,结合Dice损失和BCE损失以优化最终预测结果与真实掩膜之间的重叠度及像素级一致性。c 速度-精度权衡分析(仅供参考):c部分对比了EmbedMedSAM模型在不同平台和数据集上的帧率(FPS)与Dice分数,表明该模型在推理速度与分割性能之间实现了良好平衡。

结果
图像编码器蒸馏的训练细节
步骤1:掩码图像预训练。 在此步骤中,我们将RepViT作为图像编码器集成到MedSAM框架内,同时保持提示编码器和掩码解码器不变。我们使用ImageNet-21K数据集中0.5%的数据进行掩码图像预训练,且不使用任何标签标注。预训练期间,重建目标是原始MedSAM图像编码器产生的特征图。模型使用均方误差(MSE)损失进行优化。训练共进行10个epoch,批量大小为16,采用AdamW优化器,学习率为2e-3,β₁ = 0.9,β₂ = 0.95。训练在4块NVIDIA A800 80GB GPU上进行。
步骤2:知识蒸馏训练。 在此步骤中,我们将MedSAM图像编码器的语义特征蒸馏到前一阶段预训练的轻量级编码器中。此训练在COSMOS 1050K数据集²⁰上进行。模型训练100个epoch,批量大小为16,采用Adam优化器,初始学习率为0.0001,权重衰减为0.001。为防止过拟合,我们采用早停策略,当验证损失不再下降时终止训练。此外,当验证损失趋于平稳时,学习率调度器会将学习率降低10倍。训练在4块NVIDIA A800 80GB GPU上进行。
自提示训练框架的训练细节
在此阶段,使用表1所列的源域数据集进行训练,利用其训练集和验证集进行模型优化。评估则在源域的测试集和未见过的外部数据集上进行。默认使用AdamW优化器,初始学习率为1e-3,权重衰减为0.01。学习率根据验证损失,采用ReduceLROnPlateau策略动态调整。模型训练100个epoch,批量大小为8。所有实验均可使用两块NVIDIA 2080Ti (11GB) GPU完成。
表一
表1 | 实验中使用的专有数据集详情

评估指标
为全面评估医学图像分割性能,我们采用三项标准指标:Dice分数、豪斯多夫距离(HD)和交并比(IoU)。这些指标量化预测掩码与真实标签(GT)之间的相似度。此外,我们报告模型参数量、浮点运算次数(FLOPs)和每秒帧数(FPS)以评估计算效率。
需注意,在对比实验中,所有竞争模型仅报告其在测试集上的平均最佳Dice和HD分数25。相比之下,我们提出的Embed-MedSAM模型报告的是使用不同随机种子五次运行的平均标准差,或在特定情况下仅报告均值,以更准确地反映模型稳定性和整体性能。后续实验的构建参考25和26。
数据集
在第一训练阶段,我们使用ImageNet-21K数据集进行掩码图像预训练,并使用COSMOS 1050K医学影像数据集进行知识蒸馏。需注意,COSMOS 1050K由多个公共医学数据集组成。其中部分数据集需获得官方授权方可使用。因此,本研究中我们仅使用两类数据集:(1) 已成功获得使用许可的公开可用数据集;(2) 根据原始许可证可自由使用、无需额外授权的数据集。我们严格遵守所有使用条款和伦理要求。数据仅用于研究目的,不涉及任何商业用途或违反政策的行为。此外,COSMOS 1050K中的部分公共数据集可能与表1所列的"已见"或"未见"数据集存在重叠。为确保公平评估并避免潜在的数据泄露,我们采取了两项关键措施:1. 在Embed-MedSAM训练期间,我们仅使用COSMOS 1050K子数据集中的训练集和验证集。所有测试集均被严格排除。2. 我们使用了COSMOS 1050K的一个过滤子集。我们的团队对照COSMOS 1050K列表,手动检查了每个未见数据集。任何重叠或可能重叠的数据集均已从训练数据中完全移除。通过这两项措施,我们有效避免了COSMOS 1050K与用于内部和外部评估的测试数据集之间的任何重叠。这确保了模型评估结果的客观性和可靠性。
在第二训练阶段,我们引入了表1所列的16个下游任务数据集。这些数据集分为两类:源域(内部数据集)和未见域(外部数据集)。内部数据集依据其官方划分,分为训练集、验证集和测试集。训练集和验证集用于训练基线模型和我们的Embed-MedSAM。外部数据集仅用于推理,不参与任何训练,用于评估模型在未见域上的泛化能力。
有监督分割性能评估
为验证所提轻量级Embed-MedSAM框架在有监督医学语义分割任务中的有效性,我们与主流全自动分割架构及若干轻量级SAM变体进行了系统比较。这些对比实验在表1所列的六种医学成像模态上进行。需注意,尽管原始SAM和MedSAM等模型具有强大的泛化能力,但它们被排除在此比较之外。其极其庞大的参数量------比轻量级SAM模型大60倍以上------使得它们在资源受限设备上的部署不切实际。
关于训练策略,所有支持提示输入的轻量级SAM模型,均从官方发布的预训练权重开始,在每个训练数据集上以有监督方式进一步训练。训练采用单点提示模式,每个提示点通过使用OpenCV中的ConnectedComponentsWithStats函数计算每个实例的质心生成。此外,这些模型还支持自动掩码生成模式(用✗表示),用于评估其在无提示条件下的性能。
所有对比模型的训练策略------包括网络架构、超参数设置和训练流程------均严格遵循其原始出版物中的描述。用于评估的数据集包括表1中的S1、S2、S3、S6、S7和S9。每个数据集依据其官方协议划分为训练集、验证集和测试集。有监督训练后在测试集上的最终结果报告于表2。
主流全自动分割架构。 如表所示,尽管主流全自动分割架构在参数量上与轻量级SAM变体相当,但其推理成本(以FLOPs衡量)明显更高。例如,MADGNet需要210.65 GFLOPs,而所提模型仅消耗33.5 GFLOPs,前者计算开销是后者的6倍。这凸显了传统架构在实际部署中的低效性,尤其是在资源有限的平台上。在分割精度方面,所提模型在S1、S2、S6和S7数据集上取得了最佳结果,在S3和S9上接近次优。总体而言,与现有主流全自动分割方法相比,Embed-MedSAM模型在轻量化设计与分割精度之间实现了最佳权衡。这使其特别适合部署在资源受限的边缘设备或移动平台上的医学图像分割任务中。
无提示轻量级SAM系列模型。 同时,我们对Embed-MedSAM与其他轻量级SAM变体进行了全面比较。在无提示设置下,除ESP-MedSAM外,所提模型在所有数据集上的分割精度均显著优于所有竞争模型。尽管ESP-MedSAM在多数数据集上的精度略高于我们的模型,但性能差距微乎其微。例如,在S1数据集上,ESP-MedSAM达到88.52%,仅比我们的模型(88.39%)高0.13%。类似地,在S2数据集上,差距仅为0.11%。然而,在模型效率方面,我们的方法具有显著优势。ESP-MedSAM拥有28.46M参数------几乎是我们的模型(10.6M)的三倍。其推理复杂度也显著更高,为55.88 GFLOPs,而我们的模型仅33.5 GFLOPs。
此外,如表8所示,我们进一步验证了Embed-MedSAM的轻量化优势。尽管ESP-MedSAM与Embed-MedSAM分割精度相当,但ESP-MedSAM因模型体积大、推理复杂度高,无法在iPhone 14等小型嵌入式设备上运行。相比之下,Embed-MedSAM可成功部署在iPhone 14上,达到近30 FPS的实时推理速度。这些结果表明,Embed-MedSAM不仅保持了高分割精度,还提供了卓越的部署效率,非常适用于移动和资源受限设备上的医学图像分割任务。
此外,我们在所有数据集上进行了Wilcoxon符号秩检验,以比较Embed-MedSAM与其他基线模型。为此,我们对每个基线模型和Embed-MedSAM分别进行了五次额外的独立实验(后续实验均如此处理)。结果显示存在统计学显著差异(p < 0.05),如表p值行的"ᵃ"所示,进一步证实了我们所提方法在精度和鲁棒性上的优越性。
图4展示了若干主流模型在六个已见域数据集上的可视化分割结果。从图中可观察到,在多数基准数据集上,ESP-MedSAM与所提Embed-MedSAM的分割结果相似,且均接近真实标签。在S6和S7数据集上,Embed-MedSAM的分割结果相比ESP-MedSAM更接近真实标签,显示出更好的边界捕捉和细节保持能力。同样值得注意的是,Embed-MedSAM仅使用ESP-MedSAM三分之一的参数和一半的推理复杂度,就达到了相当甚至更优的分割性能,进一步证明了其在分割精度与计算效率之间的良好平衡。
点提示轻量级SAM系列模型。 最后,我们通过引入点提示,进一步评估了轻量级可提示SAM模型。结果表明,与无提示设置相比,点提示模式------尽管更耗费人力------确实提升了这些模型的分割精度。然而,这种提升是以增加对终端用户的专业要求为代价的。而且,即使使用点提示,所有轻量级SAM模型在任何数据集上的性能均未超越我们的无提示模型。这表明,我们的方法不仅消除了对人工交互的需求------从而增强了易用性和可及性------还实现了更优的分割性能。这些发现进一步凸显了Embed-MedSAM在实际医学影像应用中的实用性和潜力。同样需注意,ESP-MedSAM未包含在此实验中,因为它不支持基于提示的输入。
基于扩散的医学图像分割框架架构。 为进一步验证所提轻量级Embed-MedSAM模型在医学图像分割任务中的整体性能,我们将近期性能瞩目的基于扩散的分割框架作为额外基线。在表1所列的S4、S5、S8和S10数据集上进行了广泛评估,结果汇总于表3。
如表所示,Embed-MedSAM的分割精度略低于MedSegDiff-V2,但性能差距极小。例如,在S4数据集上,其Dice分数仅低0.4%。值得注意的是,MedSegDiff-V2的推理复杂度高达983 GFLOPs------比Embed-MedSAM(33.5 GFLOPs)高出近30倍。这种极高的计算需求严重限制了扩散模型在实际应用中的实用性。相比之下,Embed-MedSAM以低得多的计算成本保持了有竞争力的精度,从而在效率与性能之间实现了更优的权衡。
我们进一步在两个大规模医学图像分割基准(包括BTCV)上评估了Embed-MedSAM,结果如表4所示。与表3结果一致,Embed-MedSAM的分割精度与MedSegDiff-V2相当,而其推理复杂度仅为后者的1/30。这进一步证实了Embed-MedSAM在资源受限环境中部署的实际适用性。
对未见目标域的泛化能力
为进一步评估所提模型在未见目标域上的泛化能力,我们设计了一个跨数据集泛化评估实验。具体而言,我们复用表2实验中的模型,确保它们在训练的任一阶段均未接触过六个目标数据集------T1、T2、T3、T6、T7和T9。此设置构建了一个严格的未见域测试场景(外部数据集)。例如,表2中在S1数据集上训练的模型直接在对应的T1数据集上进行评估,从而检验其跨不同数据源和领域偏移的泛化能力。
最终评估结果见表5。如表所示,结果与表2报告的结果一致。具体地,Embed-MedSAM在所有未见数据集上均优于所有其他无提示轻量级SAM模型,ESP-MedSAM除外。尽管ESP-MedSAM在少数数据集上精度略高,但提升幅度微乎其微。更重要的是,其计算成本显著更高------其参数量约为我们模型的三倍,FLOPs接近两倍。此外,ESP-MedSAM无法部署在iPhone 14等小型嵌入式设备上,严重限制了其实际应用性。
对于可提示SAM模型,尽管使用点提示带来了适度的精度提升,但在任何未见数据集上,其性能均未超越我们的无提示Embed-MedSAM。这表明,Embed-MedSAM在不依赖人工输入的情况下,展现出更强的泛化性和鲁棒性。
此外,图5展示了若干模型在各种未见域任务上的分割结果。可视化对比清晰表明,多数模型应用于未见数据时泛化能力有限,导致性能显著下降。然而,ESP-MedSAM和所提Embed-MedSAM仍显示出与图4相似的结果,分割输出接近真实标签。考虑到Embed-MedSAM仅使用ESP-MedSAM一半的参数和三分之一的推理复杂度,却达到了相当甚至更优的视觉分割性能,这进一步证实了其在未见域任务中分割精度与计算效率之间的卓越平衡。
消融研究
关键组件。 本节我们设计了一系列消融实验,聚焦于Embed-MedSAM的三个关键模块。目标是评估每个模块对模型性能的独立贡献。最终结果见表6。需注意,表中的EKD(编码器知识蒸馏)不仅指蒸馏过程本身,还意味着图像编码器已被替换为轻量级结构。因此,从实验2到实验5,EKD均标记为启用。
具体地,与原始MedSAM基线(实验1)相比,完整版Embed-MedSAM(实验5)在已见和未见数据集上均取得了显著的性能提升。在已见数据集上,平均Dice分数从85.67%提升至89.64%,豪斯多夫距离(HD)从97.56降至79.04。在未见数据集上,Dice分数从72.29%提升至75.86%,HD从179.66降至146.82。尽管有这些提升,Embed-MedSAM的参数量和FLOPs相比MedSAM均大幅减少,展现出轻量化设计的强大优势。
在编码器蒸馏阶段,掩码图像预训练(MIP)模块对增强模型的表示能力和精度起到了积极作用。例如,在实验3中,引入MIP后,已见数据集上的Dice分数从82.73%(实验2)提升至84.36%,HD也明显降低。这表明了MIP在特征学习早期阶段的重要性。
在此实验中,MedSAM基线在推理时不使用任何外部提示(如点或框),这限制了其完整性能。一旦启用自提示训练(SPT)模块,模型便不再依赖人工提示。相反,它从初始预测生成伪提示,并用其引导更准确的第二次预测。结果表明,在所有启用SPT的配置中,Dice分数均显著提升。这证实了自提示机制极大地增强了模型在无提示条件下的适应性和易用性。
掩码图像预训练的Epoch数。 默认设置下,我们将图像编码器蒸馏(IED)阶段中掩码图像预训练(MIP)的训练epoch数设为不少于10。为验证此选择,我们进行了消融研究。如图6所示,随着训练epoch数从1增至10,Embed-MedSAM在已见和未见域上的分割精度先显著提升,后逐渐趋于稳定。这表明,在通用域图像上进行适度的预训练,可能有助于模型学习具有良好泛化性和迁移性的结构表示。
然而,当训练epoch数进一步增至11至15时,模型性能开始趋于平稳。在某些任务中,观察到轻微下降。我们推测,过度训练可能导致轻量级编码器过拟合通用域的视觉模式。这可能降低其捕获医学图像结构特征的能力,从而影响在下游分割任务上的泛化性。
在本研究的其余训练阶段,我们主要采用早停策略动态控制训练过程。不同任务的实际训练epoch数各异。因此,无需对这些阶段的固定epoch设置进行额外的消融研究。
损失权重敏感性分析。 本研究中,我们为EKD(编码器知识蒸馏)和SPT(自提示训练)设计了复合损失函数。每个损失函数整合了多个监督项,并引入了相应的损失权重因子(λ₁至λ₄)以平衡不同子任务的贡献。为评估各种权重配置对模型性能的影响,我们进行了消融研究,结果见表7。
如表7所示,各损失成分的不同权重设置对已见和未见域的Dice分数有显著影响。对于语义一致性损失(λ₁),当权重设为5 × 10⁻⁵时性能最佳,已见域和未见域的Dice分数分别为89.64和75.86。这表明适度的KL散度约束有助于在自提示训练中稳定语义对齐,提升泛化性。对于几何对齐损失(λ₂),最佳设置为λ₂ = 0.1,证实了监督伪掩码与真实框之间的对齐可改善分割质量。在分割精度方面,当Dice损失权重(λ₃)设为0.8、BCE损失权重(λ₄)设为0.2时,性能最佳。当这些权重过小或过大时,模型性能均会下降,表明适当的权重选择对稳定优化至关重要。
总体而言,默认设置在四项损失之间取得了良好平衡。每项损失从不同角度提供监督:语义一致性(L_KL)、几何对齐(L_Box)以及像素级精度(L_Dice和L_BCE)。这种多层级监督有助于模型在多样的医学图像数据集上实现良好的泛化和分割性能。
需澄清,上述默认权重设置代表了在所有数据集上整体性能最佳的配置。然而,在实际应用中,这些权重应根据具体训练反馈动态调整,以进一步优化模型在目标下游任务上的性能。
嵌入式部署
在先前的实验中,模型的训练和评估主要在NVIDIA 2080Ti高性能GPU平台上进行。然而,鉴于所提Embed-MedSAM专为资源受限边缘设备设计,其在嵌入式平台上的可行性尤为重要。因此,本节我们在低算力移动平台------iPhone 14上对模型的推理效率进行了实证评估。具体而言,我们遵循EdgeSAM论文中描述的部署方法。模型通过Core ML工具编译部署至iPhone 14。模型吞吐量使用Xcode提供的基准测试工具测量。最终结果见表8。
需注意,部署细节和硬件差异会导致结果显著不同。为确保公平比较,所有基线模型我们均采用其原始论文中报告的FPS值,而非本地复现所得值。
如表所示,原始SAM、MedSAM及效率较低的变体EfficientSAM-S和ESP-MedSAM,均在iPhone 14(本地独立测试)上遇到内存不足(OOM)问题。这表明这些模型因参数量大、计算需求高,不适合部署在内存有限的移动或嵌入式设备上。在可成功运行的轻量级SAM变体中,我们观察到其在未见域上的分割性能明显低于所提Embed-MedSAM。尽管MobileSAM和RepViT-SAM等模型满足低资源占用限制,但它们的泛化能力差,Dice分数低,使其在实际医学分割任务中可靠性不足。
相比之下,Embed-MedSAM可高效部署在iPhone 14等资源有限的设备上,同时在未见域任务中保持高分割精度。这表明我们的模型在轻量化设计与鲁棒性能之间取得了良好平衡,非常适合在边缘设备上进行实际医学图像分析。
此外,我们报告了Embed-MedSAM在iPhone 14上手动操作并执行连续推理一小时期间的运行时资源使用情况。模型消耗约25%的电量,峰值内存使用达1.9GB(超出),CPU使用率维持约30%(超出)。设备表面温度保持稳定,仅感知轻微升温。这些结果进一步表明,Embed-MedSAM能够在小型嵌入式设备上以低资源消耗稳定运行,证实了其实际边缘部署的潜力。
讨论
本文提出了Embed-MedSAM,一个专为现实世界中资源受限环境(如冲突地区、农村地区及欠发达区域)下的医学影像应用量身定制的轻量级、无提示分割模型。
Embed-MedSAM解决了现有基于SAM的模型的两大主要局限。首先,当前许多模型需要大量内存和计算资源。为克服此问题,我们将原始MedSAM编码器替换为轻量级RepViT主干网络,并采用两阶段蒸馏策略。编码器首先在ImageNet-21K上使用掩码图像重建进行预训练,随后在COSMOS 1050K医学数据集上使用MSE损失进行微调。这使得轻量级编码器能够保留MedSAM的语义表示能力,同时显著减小模型体积和推理延迟。
其次,我们通过引入自提示机制,完全消除了推理过程中对人工提示的依赖。在没有任何提示输入的情况下,模型利用第一次前向传播生成的分割掩码,通过其最小外接矩形推导出伪框提示。随后,此伪提示由提示编码器精炼,以引导第二次更准确的预测。为抑制不准确伪提示带来的误差累积,我们在训练中引入KL散度损失,强制两次预测之间的语义一致性。此外,在伪框与真实掩码框之间施加L1损失以增强几何对齐。在掩码层面,结合Dice损失和二元交叉熵(BCE)损失,优化预测掩码与真实掩码之间的重叠度和像素级精度。该框架在训练和推理阶段均消除了对手动提示的依赖,确保了强大的泛化性和部署效率。
跨17个数据集、涵盖7种以上医学成像模态的实验结果表明,Embed-MedSAM在源(已见)域和目标(未见)域均展现出强大性能。该模型在分割精度与部署效率之间提供了良好的权衡,并能在iPhone 14等设备上以近30 FPS的速度运行。这些结果表明,通用医学图像分割无需依赖高性能硬件和人工输入是可行的。
尽管Embed-MedSAM在推理时轻量高效,但其初始训练仍需一定水平的计算资源。具体而言,第一阶段编码器蒸馏采用两步策略,需在四块NVIDIA A800 GPU上进行超过一天的训练。相比之下,第二阶段自提示引导掩码解码器训练专注于使模型适应特定下游任务,训练成本显著降低。例如,在REFUGE数据集上,此阶段仅需约30分钟,使用两块NVIDIA 2080Ti 11GB GPU即可完成。因此,在一次性完成第一阶段蒸馏训练后,用户可轻松使用低成本硬件针对自己的医学图像分割任务对模型进行微调。这种范式极大增强了Embed-MedSAM在实际医学应用中的实用性和可扩展性。
总之,Embed-MedSAM为低资源环境下的医学图像分割提供了一个高效且可及的基础模型。其轻量化设计、无提示推理及鲁棒泛化能力,使其非常适合在基础设施和专业知识有限的真实临床场景中部署。通过一次性训练投入,该模型支持广泛部署,并有望改善服务不足地区的诊断可及性和效率。
方法
MedSAM基础
SAM是一个高度通用的图像分割框架,能以最少的人工输入处理多种任务。基于SAM,MedSAM²¹通过在大规模医学影像数据集上训练,将其能力扩展到医学领域。MedSAM架构包含三个主要组件:图像编码器、提示编码器和掩码解码器(见图1)。
图像编码器。 图像编码器是MedSAM中计算最密集的组件。它采用基于视觉Transformer(ViTDet)²⁷的层次化架构,使模型能够捕获多尺度特征。该编码器接收固定尺寸(3, 1024, 1024)的输入,输出尺寸为(256, 64, 64)的特征图。这一强大的主干网络使MedSAM能够有效提取深层视觉特征,支持跨多种图像域的鲁棒性能。
提示编码器。 MedSAM支持多种类型的提示,包括点、框、掩码和文本,以引导准确的掩码生成。提示编码器处理这些输入以增强分割。稀疏提示(如点和框)使用随机位置嵌入以及类型指示标记进行编码,以区分提示类型。对于更复杂的提示(如掩码和自由格式文本),MedSAM分别使用小型CNN和CLIP文本编码器²⁸。本研究中,我们关注交互式分割任务中常用的点和框提示。
掩码解码器。 掩码解码器生成最终的分割输出。它接收来自编码器的图像特征、嵌入的提示(点或框)、掩码预测标记以及IoU(交并比)预测标记作为输入。这些输入通过一个轻量级双向Transformer处理。每个掩码标记被转换为一个动态线性分类器,用于计算每个空间位置的前景概率。IoU标记提供预测掩码的置信度分数。此架构使MedSAM即使在挑战性条件下也能生成准确且自适应的掩码。
大规模医学图像分割数据集。 MedSAM强大性能背后的一个主要因素是其在大规模医学图像数据集上的训练。该数据集包含1,570,263对图像-掩码,涵盖10种成像模态和超过30种癌症类型。作为迄今为止最大的分割数据集,它显著提升了MedSAM在多种医学任务上的零样本泛化能力。
研究目标
本研究的主要目标是基于先进的MedSAM框架开发一个医学图像分割模型,需达成两个关键目标:能够在资源受限的嵌入式设备上部署,以及在推理时消除对提示输入的需求。模型必须显著降低推理时的计算成本,同时保持原始MedSAM所达到的高分割精度。需注意,MedSAM中图像编码器包含超过6亿参数,而提示编码器和掩码解码器的参数总数不到400万。因此,本研究应对两大挑战:(1) 如何用轻量级替代方案替换大规模图像编码器,同时保持MedSAM级别的分割精度,实现在边缘设备上的高效部署;(2) 如何在推理时消除对提示输入的依赖。
轻量级MedSAM
我们引入了基于RepViT²⁴的轻量级图像编码器。RepViT是一种源自MobileNetV3²⁹的卷积神经网络(CNN),并融合了现代轻量级视觉Transformer(ViTs)的设计优势。与同类规模的其他轻量级ViT模型相比,它实现了更优的性能,并在iPhone 12等移动设备上展现出更低的延迟。
本工作中,我们选择RepViT-M1.0变体,其CPU推理速度显著快于使用TinyViT主干网络³⁰的基线模型。我们编码器的宏观结构如图3所示,包含一个茎干层、两个下采样层和一个编码层。在这些阶段之间,我们按1:1:8的比例插入多个RepViT块。RepViT还采用结构重参数化技术³¹,减少了计算和内存开销,从而在不牺牲精度的情况下,在资源受限设备上实现更快推理。
为确保轻量级MedSAM达到与原始MedSAM相当的性能,同时消除其对提示输入的依赖,我们采用了两阶段训练策略,其灵感来源于Zhang等人³²提出的"分而治之训练范式"。具体地,我们将训练过程分解为两个独立的子任务:图像编码器蒸馏和提示引导掩码解码器训练。
图像编码器蒸馏
将原始大规模图像编码器替换为轻量级对应物后,常规方法通常在目标域(如医学影像数据)上直接进行知识蒸馏,以迁移原始编码器的特征提取和语义理解能力。然而,由于轻量级模型容量有限,这种方法容易导致过拟合。
为缓解此问题,我们提出一种新颖的两步蒸馏策略,如算法1所示。第一步,在部分ImageNet-21K数据集33上进行掩码图像预训练,利用MedSAM的重建监督迁移通用视觉知识。第二步,在大规模医学数据集COSMOS 1050K的子集上应用均方误差(MSE)损失的优化23。这引导轻量级编码器与MedSAM的语义特征对齐,有效继承其在医学图像分割任务中的领域适应性。
掩码图像预训练
在此微调阶段,受MAE和EfficientSAM³³启发,我们引入了一种MedSAM引导的掩码图像预训练策略。如图2a步骤1所示,轻量级编码器从未掩码令牌中提取潜在特征,而交叉注意力解码器利用未掩码特征和MedSAM提取的目标特征重建掩码令牌。这使得编码器能够在ImageNet-21K上学习MedSAM风格的表示,无需任何医学标注。
算法1. 两步蒸馏策略流程

交叉注意力解码器。 仅重建掩码令牌,未掩码令牌的编码器输出作为参考键和值。掩码令牌被视为查询。然后将掩码和未掩码令牌的输出合并并重排序,生成最终的潜在嵌入。
线性投影头。 为匹配MedSAM特征的维度,我们对解码器输出应用轻量级线性投影头,以实现特征级对齐。
重建损失。 在每个训练步骤中,过程包括:通过MedSAM图像编码器的前向传播提取特征,然后在MAE模型中进行前向和后向传播。随后,通过比较MedSAM编码器的输出与MAE中线性投影头的输出来计算重建损失。
令MedSAM图像编码器表示为 f_sam,MAE的编码器和解码器分别为 g_e 和 g_d,权重为 W_e 和 W_d。将线性投影头定义为 hθ,权重为 Wθ。假设输入令牌由 {x_i}^N_{i=1} 表示,其中 N 是令牌总数。这些令牌根据预定义的掩码比例分为未掩码令牌 {x_i}{i∈U} 和掩码令牌 {x_i}{i∈M}。用 ϕ 表示重排序函数,⊕ 表示合并函数。
MedSAM编码器输出的目标特征可表示为 f_sam(x) = f_sam({x_i}^N_{i=1})。MAE编码器对未掩码令牌的输出为 g_e({x_i}{i∈U}),而解码器生成 g_d({x_i}{i∈M})。线性投影头的输出则为:

因此,我们的目标重建损失可表示为:

其中 N 是输入令牌数,||·|| 表示范数。实验中我们使用 l₂ 范数作为重建损失。通过最小化重建损失 L**{W_e, W_d, W**θ},我们的编码器 g_e 被优化为作为图像骨干网络,像MedSAM图像编码器一样提取特征。我们的编码器、解码器和线性投影头被优化以从MedSAM图像编码器学习上下文建模能力。对所有令牌优化重建损失,迁移了嵌入在MedSAM中的知识。
编码器知识蒸馏
在此阶段,我们使用大规模医学影像数据集COSMOS 1050K的子集,在MedSAM图像编码器的引导下进行语义蒸馏。这有助于轻量级模型继承MedSAM在医学图像中的语义表示和结构感知能力,如图2a步骤2所示。目标是在医学领域实现有效的知识迁移和领域适配。
具体地,我们应用均方误差(MSE)损失在像素级别对齐教师模型和学生模型的特征图。此损失帮助学生模型学习教师模型的中间表示,从而改善其语义理解。
蒸馏损失定义为:

其中 f^T(x_i) 和 f^S(x_i) 分别表示来自MedSAM图像编码器(教师)和轻量级RepViT图像编码器(学生)的特征图。N 是特征图中的总像素数。此损失确保学生模型在不增加复杂度的情况下有效模仿教师模型的特征表示,使其适合部署在资源受限环境中。
自提示训练框架
为增强解码器与先前蒸馏所得轻量级图像编码器的兼容性,并消除对人工提示的依赖,我们在下游医学图像数据集上提出了一种自提示引导的解码器训练框架,受22启发。
传统模型如SAM和MedSAM在训练和推理过程中依赖外部提供的提示(如框、点或掩码)。为消除此依赖,我们引入了一个即插即用的自提示模块,如图2b所示。
令 x 表示输入图像,g_e(x) 为图像编码器提取的图像特征。在初始训练步骤中,解码器使用内部初始化的默认提示嵌入 p^(0) 生成粗略分割掩码 ŷ^(0):

接着,我们提取 ŷ^(0) 的边界生成自推导的框提示 b^(1),该提示由提示编码器 g_p 编码为更新后的提示嵌入:

然后生成精炼预测:

为确保在没有人工输入的情况下实现鲁棒的自提示学习和准确分割,我们采用了多项损失函数,监督训练过程的不同方面:
语义一致性。 为防止自生成提示带来的误差累积,我们引入Kullback-Leibler (KL)散度损失,强制初始粗略预测与精炼预测之间的语义对齐:

几何对齐。 为改善伪框提示的质量,我们使用伪掩码边界框与真实掩码边界框之间的L1损失:

分割精度。 为引导最终预测实现准确区域描绘,我们联合应用Dice损失和二元交叉熵(BCE)损失于精炼掩码预测与真实掩码之间:

总训练损失公式化为所有成分的加权和:

其中 λ₁ 至 λ₄ 是经验选择用于平衡一致性、对齐和准确性的损失权重。默认情况下,损失权重参数设为:λ₁ = 5 × 10⁻⁵,λ₂ = 0.1,λ₃ = 1,λ₄ = 1。然而,这些超参数可根据特定下游医学影像数据集的特点进行调整,以获得更好的训练性能。
数据可用性
ImageNet-21K数据集可从 https://huggingface.co/datasets/gmongaras/Imagenet21K 下载。COSMOS 1050K数据集可通过 https://github.com/yuhoo0302/Segment-Anything-Model-for-Medical-Images 提供的下载链接访问。
代码可用性
图像编码器蒸馏的步骤1代码基于EfficientSAM中的描述。图像编码器蒸馏的步骤2代码是Rep-MedSAM的实现。自提示训练框架代码参考SAM-SP,一个简单示例代码提供于 https://github.com/XinLianSJ/Embed-MedSAM。