[论文解读] ASPS: Augmented Segment Anything Model for Polyp Segmentation

Huiqian Li, Dingwen Zhang+4 more|arXiv (Cornell University)|Jun 30, 2024

Vehicle License Plate Recognition被引用 4

一句话总结

本文提出ASPS，一种针对息肉分割的新型分割一切模型（SAM）改进方法，有效缓解了领域偏移问题并提升了分布外泛化能力。通过将可训练CNN编码器与SAM的冻结ViT通过跨分支特征增强（CFA）融合，并应用不确定性引导的预测正则化（UPR）以改善置信度校准，ASPS在无需提示的情况下，于五个息肉数据集上实现了最先进性能。

ABSTRACT

Polyp segmentation plays a pivotal role in colorectal cancer diagnosis. Recently, the emergence of the Segment Anything Model (SAM) has introduced unprecedented potential for polyp segmentation, leveraging its powerful pre-training capability on large-scale datasets. However, due to the domain gap between natural and endoscopy images, SAM encounters two limitations in achieving effective performance in polyp segmentation. Firstly, its Transformer-based structure prioritizes global and low-frequency information, potentially overlooking local details, and introducing bias into the learned features. Secondly, when applied to endoscopy images, its poor out-of-distribution (OOD) performance results in substandard predictions and biased confidence output. To tackle these challenges, we introduce a novel approach named Augmented SAM for Polyp Segmentation (ASPS), equipped with two modules: Cross-branch Feature Augmentation (CFA) and Uncertainty-guided Prediction Regularization (UPR). CFA integrates a trainable CNN encoder branch with a frozen ViT encoder, enabling the integration of domain-specific knowledge while enhancing local features and high-frequency details. Moreover, UPR ingeniously leverages SAM's IoU score to mitigate uncertainty during the training procedure, thereby improving OOD performance and domain generalization. Extensive experimental results demonstrate the effectiveness and utility of the proposed method in improving SAM's performance in polyp segmentation. Our code is available at https://github.com/HuiqianLi/ASPS.

研究动机与目标

• 解决SAM预训练中自然图像与内镜息肉图像之间的领域差距，该差距限制了性能表现。
• 提升SAM在息肉图像中捕捉局部和高频特征的能力，这些特征常因模型对全局、低频模式的关注而被忽略。
• 降低预测不确定性，改善对分布外（OOD）内镜图像的置信度校准。
• 通过端到端、无需提示的训练，消除临床应用中对提示的依赖。
• 通过轻量化、高效的适配框架，提升领域泛化能力和鲁棒性。

提出的方法

• 提出跨分支特征增强（CFA），将可训练CNN编码器的特征与冻结的视觉Transformer（ViT）编码器特征融合，以增强多尺度和高频特征学习能力。
• 采用不确定性引导的预测正则化（UPR），利用真实标签作为提示来校准置信度分数并减少训练过程中的不确定性。
• 对颈部和Transformer模块中的归一化层进行适配，以提升对内镜领域分布的适应能力。
• 利用傅里叶分析验证CNN分支相比ViT基线能捕捉到更多高频内容。
• 实现端到端无提示训练，使模型可直接对内镜图像进行推理。
• 在UPR中利用SAM的IoU预测作为监督信号，以指导不确定性降低并提升分布外泛化能力。

实验结果

研究问题

• RQ1 与纯ViT-based SAM相比，混合CNN-ViT编码器架构是否能提升息肉分割的特征表示能力？
• RQ2 在训练中引入不确定性正则化是否能增强分布外泛化能力并改善置信度校准？
• RQ3 能否在不微调原始SAM权重或依赖提示的前提下，实现领域特定的特征增强？
• RQ4 多级特征融合与位置嵌入替换的结合如何影响分割精度和鲁棒性？
• RQ5 所提方法在多样化息肉数据集上相较于现有SAM-based模型的性能优势有多大？

主要发现

• ASPS在五个息肉数据集上平均Dice得分为0.914，平均IoU为0.843，优于基线SAM及MedSAM、SurgicalSAM等先前方法。
• 仅CFA模块相比ViT-B基线，平均Dice提升31.7%，平均IoU提升41.4%。
• CFA与UPR联合使用在CVC-ClinicDB数据集上达到最佳性能，Dice得分为0.950，IoU为0.905。
• 采用颈部层归一化（NN）与提示引导正则化的UPR表现最佳，证明其在不确定性减少方面的有效性。
• 傅里叶分析证实，CNN分支相比ViT编码器捕捉到显著更多的高频内容，支持了局部细节学习能力的提升。
• 定性结果表明，ASPS的预测结果更接近真实标签，尤其在捕捉息肉的精细边界和纹理细节方面表现更优。