[论文解读] Customized Segment Anything Model for Medical Image Segmentation

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

Advanced Neural Network Applications被引用 121

一句话总结

SAMed 自定义 Segment Anything Model (SAM) 用于医用图像分割，采用 LoRA 微调，在最小部署/存储开销下实现具竞争力的结果。

ABSTRACT

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

研究动机与目标

• 将 SAM 扩展到具有有意义组织标签的语义医用图像分割。
• 通过仅更新 SAM 参数的一小部分，实现高效微调。
• 利用 warmup 和 AdamW 稳定训练并改善收敛。
• 在 Synapse 多器官数据集上展示具竞争力的性能。
• 展示 SAMed 在保持 SAM 兼容性的同时，部署/存储开销只是边际。

提出的方法

• 冻结 SAM 图像编码器，并对其变换器块应用 LoRA 以学习医学特征。
• 为语义分割微调提示编码器和掩码解码器（可选地使用 LoRA）。
• 将 SAM 的输出适配为预测与 k 个类别（包括背景）对应的 k 个语义掩码，并通过在类别维度上对 Softmax 和 ArgMax 计算最终 S。
• 使用交叉熵和 Dice 损失，并对地面真值进行下采样以进行训练监督。
• 采用 warmup 阶段和 AdamW 优化器以稳定训练并改善收敛。
• 证明与 SAM 的兼容性并展示更新参数数量的下降（例如仅在图像编码器上使用 LoRA 时为 18.81M）。

实验结果

研究问题

• RQ1 在对医用数据进行 LoRA 微调后，SAM 是否能够对医用图像执行语义分割？
• RQ2 更新 SAM 参数的一小部分是否在降低部署/存储开销的同时获得具竞争力的分割精度？
• RQ3 哪些训练策略（warmup、AdamW）能改善医用数据上 SAM 调整的收敛和性能？
• RQ4 SAMed 与 Synapse 上的最先进医用分割模型在 DSC 和 HD 上的对比如何？
• RQ5 在保持 SAM 兼容性的同时，SAMed 是否能够为不同组织提供有意义的语义标签？

主要发现

• SAMed 在 Synapse 多器官数据集上达到 81.88 DSC 和 20.64 HD，与最先进基线并驾齐驱。
• SAMed 在 Synapse 结果中胰腺和胃分割方面达到最前沿的性能。
• 仅更新极少数 SAM 参数（例如 18.81M 相较于原始 358M；原始大小的 5.25%），使部署/存储开销微小。
• 对图像编码器应用 LoRA（和可选的掩码解码器）比仅更新掩码解码器获得更好性能。
• Warmup 和 AdamW 优化器显著稳定训练、改善收敛和最终损失。
• SAMed 仍然与 SAM 完全兼容，可作为插件来提升 SAM 在医用图像分割中的应用。