【论文阅读笔记】ASPS: Augmented Segment Anything Model for Polyp Segmentation

1.论文介绍

ASPS: Augmented Segment Anything Model for Polyp Segmentation
ASPS：用于息肉分割的扩展SAM模型
2024年 arxiv
Paper Code

2.摘要

息肉分割在结直肠癌诊断中起着至关重要的作用。最近，Segment Anything Model(SAM)的出现利用其在大规模数据集上的强大预训练能力，为息肉分割带来了前所未有的潜力。然而，由于自然图像和内窥镜图像之间的区域差距，SAM在实现有效的息肉分割方面遇到了两个限制。首先，它的基于变压器的结构优先考虑全局和低频信息，可能会忽略局部细节，并在学习的特征中引入偏差。其次，当应用于内窥镜图像时，其较差的分布(OOD)性能导致不符合标准的预测和有偏差的置信度输出。为了应对这些挑战，我们提出了一种新的用于息肉分割的扩展SAM(ASPS)方法，该方法包括两个模块：跨分支特征增强(CFA)和不确定引导预测正则化(UPR)。CFA将可训练的CNN编码器分支与冻结的VIT编码器集成在一起，在增强局部特征和高频细节的同时，实现了特定领域知识的整合。此外，UPR巧妙地利用SAM的IOU分数来减少训练过程中的不确定性，从而提高了面向对象设计的性能和领域泛化。大量的实验结果证明了该方法在提高SAM在息肉分割中的性能方面的有效性和实用性。

Keywords： SAM，CNN+ViT，跨分支特征增强，不确定引导预测正则化

3.Introduction

自动息肉分割是诊断结直肠癌的关键工具，有助于有效的干预和及时的治疗策略。最近SAM被引入，由于其巨大的模型大小和数据量，这种创新的方法为息肉分割领域引入了新的视角。它还具有增强的表示和特征提取能力，超过了现有的方法。然而，由于训练数据和内窥镜图像之间的域差距，SAM在息肉分割任务中的性能并不令人满意。这导致了两个主要问题：第一，SAM不能充分捕捉息肉图像的显著特征，导致其学习表示的偏差。其次，它对分布外数据产生了错误的预测和不准确的置信度估计。此外，由于SAM依赖于提示，大大阻碍了其在临床应用中的便利性。尽管有几种方法改进了SAM，但这些方法要么依赖于提示，要么直接微调实体模型。Samus有效地集成了CNN和VIT，但其设计相当复杂，特别适合处理小图像。因此，这些方法的有效性在某种程度上受到了限制。人们已经提出了各种方法来解决语义切分中无监督领域自适应的挑战。MIC提出了一种用于目标领域情境学习的掩蔽图像一致性模型；情境感知领域自适应通过交叉注意改进了情境迁移。然而，特定领域的信息集成和不确定性减少仍然没有被探索。为了解决这些问题，本文从领域自适应的角度提出了一种新的基于SAM的方法，旨在增强特征提取能力和泛化能力，而不依赖于提示。本文提出了交叉分支特征增强模块(CFA)和不确定性引导预测正则化模块(UPR)。CFA加入了一个额外的可训练卷积神经网络(CNN)编码器分支，该分支补充了冻结视觉转换器(VIT)编码器，以捕获多尺度和多层次的特征。UPR调整归一化层以促进内窥镜领域的自适应，并利用提示确保准确的置信度估计，从而提高SAM的面向对象设计性能。

4.模型结构详解

在这里插入图片描述
CFA模块集成了CNN编码器特征和全局VIT信息，导致了广义特征表示学习。这种集成通过将深层信息聚合到浅层并结合来自浅层的位置信息来促进精细化分割输出。同时，UPR的设计是为了最大限度地减少训练期间的不确定性和校准置信度。UPR利用基于不确定性的训练策略，利用gt作为指导性“提示”。提出的网络遵循端到端训练，在没有提示的情况下，联合优化两个模块以实现最佳性能。

Cross-branch Feature Augmentation Module 跨分支特征增强模块：
SAM在息肉分割任务中仍存在一定的局限性，其中一个主要原因是SAM的图像编码器不能有效地从不可见的内窥镜图像中捕获足够的特征。为了解决这个问题，CFA模块被设计成学习多尺度特征和多层次表示，从而增强编码器的特征提取能力。

首先，为了实现自动分割，对SAM的体系结构进行了改进，去掉了它的提示输入和提示编码器部分，保留了它的图像编码和掩码解码部分。最近的研究表明，VIT更专注于低频信号，而CNN更擅长处理高频信号。因此，本文集成了一个基于CNN的并行分支来弥补高频和局部特征的缺失。此外，通过提出一个额外的多头交叉分支注意块来增强SAM的掩码解码器，以促进从VIT编码器和CNN编码器中提取的特征的整合。对于VIT分支机构的 $F_V$ 特征和CNN的 $F_C$ 特征，跨分支特征注意力可以表达如下：
在这里插入图片描述
其中 $Q=F_vW^Q$ ， $K=V=F_cW^K$ ，并且d是Fv的每个头部的通道数。考虑到CNN特征提供了更精确的位置信息，用CNN编码器的最终输出特征代替了SAM在掩码解码器中的原始位置嵌入。此外，将跨分支注意机制集成到掩码解码器的注意块中，重复此过程两次，以确保来自VIT和CNN编码器的多尺度特征的集成。
在这里插入图片描述
其次，为了获得更准确的分割结果，将来自编码器的高层上下文和低层边界信息与SAM的解码器特征相结合，以增强输出信息。具体地说，将从中间嵌入的VIT编码器获得的浅局部特征，从VIT编码器获得的最终全局特征与CNN的最终特征相结合，像上图所示的，这种方法充分利用了丰富的边缘信息、广泛的全局上下文信息和每个编码器分支的局部位置细节。因此，可以有效地集成VIT和CNN的多层次特征。

ViT，即原SAM的image encoder的最终特征和中间层特征，与CNN的最终特征相结合。并在原SAM的mask decoder中增加多头注意力，以融合ViT与CNN的特征。
首先在原本的image_encoder中提取第4、8、12层特征作为interm embedding（代码中这么设置），无prompt_encoder，在mask_decoder中，首先把来自CNN的特征，与interm embedding和image embedding融合，用CNN的特征代替ViT的位置嵌入，并在transformer块中融合CNN与image embedding特征，然后在最后把fusion的特征再加上去。
除此之外，增加一个置信度计算。

Uncertainty-guided Prediction Regularization Module 不确定性引导的预测正则化模型：

为了增强SAM的泛化能力，本文提出了一种新的训练策略，该策略包括在编码器内选择性激活LayerNorm。我们还以gt为线索，通过纠正置信度来进一步指导训练过程。由于SAM是在自然数据上训练的，其在息肉图像中的性能可能会因域转移而恶化。尽管LayerNorm的引入可能会减少训练时间，但它从根本上改变了输入数据的分布。当将SAM从自然图像转移到内窥镜图像时，数据分布和相应的特征空间分布都发生了偏移。这些分布差异可能会导致内部协变量的变化，从而影响模型的性能。为了提高SAM在内窥镜领域的泛化能力，本文对编码器的归一化层进行了微调。在这个过程中，该模型有效地适应了目标领域的数据分布，并缓解了内部协变量变化的影响。

具体地说，SAM的VIT编码器的层范数LayerNorm分为(1)transformer块norm和(2)neck层norm，如上图(A)所示。考虑到Neck层的特征更接近编码器的输出特征，最终决定训练Neck层归一化，这相当于对预先训练的VIT编码器的特征进行重新归一化。

这里提到训练策略在编码器的neck层增加LayerNorm，原代码中就是可选择的，不理解这里怎么重新处理了。

此外，以前的研究已经证明，不确定性较低的预测往往表现出优越的分布外(OOD)性能，这也有利于领域适应。SAM会生成一个IOU分数输出，这本身就代表着不确定性(或者置信度)。然而，在预测过程中，SAM可能会频繁地对看不见的数据产生高置信度的错误预测，这是不可取的。为了缓解这个问题，我们努力减少模型在训练过程中的不确定性(即增加置信度)。在[4]的启发下，我们利用gt作为提示来指导模型的学习。首先，我们将SAM的IoU得分表示为图像级别的置信度Ci。然后，我们计算像素级置信度cp，以使用公式来细化每个像素的不确定性。其中Up∈Rb×1×H×W.
在这里插入图片描述
Up表示像素不确定性，定义为 $U p = 1 - σ (∣ P ∣)$ 。这里，σ表示Sigmoid函数，而P表示输出预测。最终置信度被计算为图像级置信度和像素级置信度之和，表示为 $c = 1/2 (C i + c p)$ 。这种可信度是由伯努利分布决定的，它决定了是否将gt作为提示。换句话说，如果置信度足够低，我们认为该模型需要特定的答案提示来学习正确的掩码预测。因此，答案是必需的，作为提示，否则就没有必要。提示的权重由置信度c确定，其表示如下：
在这里插入图片描述
然而，通过最小化损失函数，模型将倾向于使c=0，从而 $P^{'}$ 将始终是GT。这意味着该模型实际上并不学习。因此，引入置信度损失来监督c，当c→为0时，置信度损失会增加，并且置信度损失的定义如下：
在这里插入图片描述
最终损失函数是分段损失Ls和置信度损失Lc之和，如公式中所定义的。5.这里，λ表示一个超参数。具体地，所采用的分段损失是CE损失、Dice损失和MSE损失的组合，如ls=Lce+0.5Ldice+Lmse。