【SAM下游任务微调】TS-SAM: Fine-Tuning Segment-Anything Model for Downstream Tasks

TS-SAM: Fine-Tuning Segment-Anything Model for Downstream Tasks

论文链接：https://arxiv.org/pdf/2408.01835

代码链接：https://github.com/maoyangou147/TS-SAM

一、摘要

  基于适配器的微调方法已经被研究用于改进SAM在下游任务上的性能。然而，经过微调的SAM与特定领域的模型之间仍然存在显著的性能差距。为了减少这种差距，作者提出了双流SAM（TS-SAM）。

  一方面，受到Parameter-Efficient Fine-Tuning（PEFT）中的侧网络灵感的启发，设计了一个轻量级的卷积侧适配器（CSA），它将SAM强大的特征集成到侧网络训练中，以实现全面的特征融合。另一方面，根据分割任务的特性，设计了多尺度细化模块（MRM）和特征融合解码器（FFD），以保持细节和语义特征。

  在三个任务领域的十个公开数据集上进行的广泛实验表明，TS-SAM不仅显著优于最近提出SAM-Adapter和SSOM，而且在性能上与最先进的特定领域模型相竞争。

二、创新点

1. 首次将侧网络引入SAM的微调中。提出了一个双流侧网络的架构，它有效地从SAM编码器中提取特征。

2. 提出了多尺度细化模块（MRM）和针对分割任务的特征融合解码器（FFD）。这些模块通过高分辨率层次特征获取细化后的目标位置信息，并在解码过程中完全融合它们，以实现详细的分割结果。

3. 对提出的TS-SAM在三个任务（包括隐身物检测COD、阴影检测和显著对象检测SOD）的十个公开数据集上进行了评估。实验结果表明，TS-SAM在这些下游任务上显著优于最近的工作，这些工作是专门为这些任务微调的SAM。甚至与专门为每个任务设计的SOTA领域特定模型相比，TS-SAM也达到了竞争性能。
   
   

三、原理

A. Overall Architecture

 图 2(a) 展示了提出的TS-SAM的整体架构。使用了预训练的SAM ViT作为骨干网络，并在下游任务中设计了一系列轻量级模块进行微调。给定一幅图像 $I\in R^{3\ast H\ast W}$，通过SAM图像编码器提取到的视觉特征 $F_{vit}\in R^{C\ast \frac{H}{16}\ast \frac{W}{16}}$ 。同时，通过堆叠的卷积侧适配器（CSA）逐层提取SAM图像编码器的特征，得到适应于下游任务的图像特征 $F_{csa}\in R^{C_1\ast \frac{H}{16}\ast \frac{W}{16}}$。为了从图像编码器提取更详细的特征，提出了多尺度细化模块（MRM）。MRM上采样图像编码器各层的特征嵌入，创建了一个层次化的特征表示 { F m r m k } k = 1 2 \{F^{k}_{mrm}\}^{2}_{k=1} {Fmrmk​}k=12​。此外，一个轻量级的门控单元连续地将SAM图像编码器的低层到高层特征融合，收集了更丰富的图像细节特征。

 解码过程没有使用SAM掩码解码器。原因在于SAM解码器需要诸如点或框之类的提示才能取得良好的效果，并且通过单一的前向传播分割多个目标是一项挑战。因此，设计了轻量级的特征融合解码器（FFD），将层次化的特征表示 { F m r m k } k = 1 2 \{F^{k}_{mrm}\}^{2}_{k=1} {Fmrmk​}k=12​注入从CSA获得的特征$F_{csa}$。这增强了特征表示，得到了一个精化的分割掩码。FFD在$F_{csa}$上进行上采样过程时，逐步融合层次化的特征表示，通过两阶段注入方法突出层次化特征表示中的关键信息。

 最后，为了减少训练成本，在训练过程中冻结SAM图像编码器。只有Convolutional Side Adapter、Multi-Scale Refinement Module和Feature Fusion Decoder这三个组件被训练，并且所有三个组件都是轻量
级的。

B. Convolutional Side Adapter

 本节详细介绍了卷积侧适配器（CSA）。受到PEFT领域适配器技术启发，对原始适配器进行了简单的修改。CSA能够有效地提取SAM图像编码器的特征，并弥合这些特征与下游任务数据之间的差距。如图2(b)所示，与原始适配器不同，CSA由两个$1\ast 1$卷积模块组成。第一个$1\ast 1$卷积模块扩展压缩的特征到SAM图像编码器的特征维度，然后与SAM图像编码器的输出特征嵌入合并。第二个$1\ast 1$卷积模块将合并的特征压缩回输入特征维度，作为下一个CSA层的输入。给定$（i-1）$-th CSA的输出特征$F_{csa}^i\in R^{C_1\ast \frac{H}{16}\ast \frac{W}{16}}$和SAM图像编码器$（j-1）$-th层的输出特征$F_{vit}^j\in R^{C\ast \frac{H}{16}\ast \frac{W}{16}}$，第$i$个CSA模块的操作可以表示如下：
$$F_{vit}^{j+1}=F_{vit}^j+con{v}_{1\ast 1}(F_{csa}^i)\text{\hspace{1em}(1)}$$

$$F_{csa}^{i+1}=con{v}_{1\ast 1}(F_{vit}^{j+1})\text{\hspace{1em}(2)}$$

其中，$con{v}_{1\ast 1}$代表$1\ast 1$卷积模块，包括一个$1\ast 1$的卷积操作，批归一化和激活函数操作。$F_{vit}^{j+1}$是SAM图像编码器的第$j$层的输入，$F^{i+1}$是第$i$个CSA的输入，CSA是轻量级，且保留适配器的精简。

C. Multi-Scale Refine Module

 在分割任务中，为了实现更精确的分割结果，模型需要能够有效地描述诸如物体边缘等细节特征。然而，在SAM图像编码器的块嵌入过程中，图像的16倍下采样可能会导致难以提取目标位置信息。因此，提出**多尺度细化模块（MRM）**来获取具有更多细节的高分辨率特征。

 图2©展示了MRM的结构。第$i$个MRM层接收来自SAM图像编码器的$(j-1)$层输出特征$F_{vit}^j$和来自MRM的第$(i-1)$层的层级特征${F_{mr{m}_i}^k}_{k=1}^2$，输出两个特征${F_{mr{m}_{i+1}}^k}_{k=1}^2$，其分辨率是$\frac{H}{8}\times \frac{W}{8}$ 和$\frac{H}{4}\times \frac{W}{4}$。首先通过一个$1\ast 1$卷积模块压缩$F_{vit}^j$的特征维度以获得${\hat{F}}_{vit}^j\in R^{C\prime \ast \frac{H}{16}\ast \frac{W}{16}}$。然后，通过一个反卷积模块处${\hat{F}}_{vit}^j$以产生更高分辨率的多尺度特征表示 { F m r m i + 1 k } k = 1 2 \{F^{k}_{mrm_{i+1}}\}^{2}_{k=1} {Fmrmi+1​k​}k=12​。上述过程可以正式化如下：
$$F_{mr{m}_{i+1}}^k=decon{v}_k(con{v}_{1\ast 1}(F_{vit}^j)),k=1,2\text{\hspace{1em}(3)}$$

其中 $deconv$ 代表反卷积模块，该模块对 ̂$F_{vit}^j$ 执行$2x$ 和 $4x$ 上采样，从而在两个不同尺度上产生高分辨率特征。进一步，获得的的高分辨率特征需要与MRM 前一层的输出层次特征合并。

 为了控制特征融合的程度，为不同尺度的高分辨率特征使用了 [10] 中提出的轻量级门控单元。该单元通过线性层和激活操作计算像素级别的权重，从而精细地控制了特征融合的程度。门控单元的操作可以形式化为以下内容：
$${\tilde{F}}_{mr{m}_{i+1}}^k=Tanh(Linear(ReLU(Linear({\hat{F}}_{mr{m}_{i+1}}^k))))\otimes {\hat{F}}_{mr{m}_{i+1}}^k,k=1,2\text{\hspace{1em}(4)}$$

 其中Linear 表示线性层，而$\otimes$表示元素相乘。最后，将这两个特征简单相加以实现特征融合：
$$F_{mr{m}_{i+1}}^k={\tilde{F}}_{mr{m}_{i+1}}^k+F_{mr{m}_i}^k,k=1,2\text{\hspace{1em}(5)}$$

D. Feature Fusion Decoder

 为了确保CSA（卷积分割注意力）和MRM（多分辨率模块）分支特征的彻底融合，提出了一种轻量级的特征融合解码器（FFD）。FFD设计用于在解码过程中将层次化特征表示注入自CSA分支获得的特点中。FFD的结构如图3所示。

 CSA分支的特征$F_{csa}$首先通过一个$1\ast 1$的卷积处理得到$F_{in,csa}$。

 对于具有$\frac{H}{8}\ast \frac{W}{8}$ 和$\frac{H}{4}\ast \frac{W}{4}$ 分辨率的层次化特征表示${F_{mrm}^k}_{k=1}^2$，目标是突出其关键特征。因此，使用池化操作对层次化特征表示进行下采样，强调其关键组成部分。这一过程可以表示如下：

$$F_{in,mrm}^k=con{v}_{1\ast 1}(F_{mrm}^k),k=1,2\text{\hspace{1em}(6)}$$

$$F_{key,mrm}^k=GAP(F_{in,mrm}^k)+GMP(F_{in,mrm}^k),k=1,2\text{\hspace{1em}(7)}$$

其中GAP（全局平均池化）和GMP （全局最大池化）分别具有$2\ast 2$的核。对于不同尺度的高分辨率特征，采用两级注入方法。首先，通过池化操作获得的关键特征$F_{key,mrm}$被注入到$F_{in,csa}$中，然后完整的、高分辨率的层次特征$F_{mrm}$被注入到$F_{in,csa}$中。以高分辨率特征$F_{mrm}^1$的注入过程为例，其分辨率是$\frac{H}{8}\ast \frac{W}{8}$ ，注入过程可以正式化如下：
$$F_{stage1,csa}=con{v}_{3\ast 3}([F_{in,csa},F_{key,mrm}^1])\text{\hspace{1em}(8)}$$
$$F_{stage1,csa}=UpSample(F_{stage1,csa})\text{\hspace{1em}(9)}$$
$$F_{stage2,csa}=con{v}_{3\ast 3}([F_{stage1,csa},F_{mrm}^1])\text{\hspace{1em}(10)}$$

其中方程（8）和方程（10）分别表示注入过程的第一阶段和第二阶段，其中表示沿着通道维度进行拼接，并且UpSample表示$2x$上采样。随后，高分辨率特征$F_{mrm}^2$的注入过程，其分辨率是$\frac{H}{4}\ast \frac{W}{4}$，与上述过程相同。通过这一两阶段的注射，层次化的特征与CSA分支的特征得到了完全整合，从而产生了一种增强的特征表示。这使得解码过程能够实现更为精细的分割结果。

四、实验

A. Datasets and Implementation

Database.
 本节报告了在SAM-Adapter[3]和SSOM [4]中使用的十种来自三个具有挑战性的下游任务的实验数据集。

1. COD: 用于测试TS-SAM 性能的四个常用数据集：CAMO [23], COD10K [24], CHAMELEON [25] 和NC4K [26]。按照 [3] 中的训练协议，将 [23] 和[24] 的训练样本合并用于模型训练，并在每个数据集上使用测试样本进行测试。

2. 暗影检测：使用了ISTD [27]数据集来测试TS-SAM的性能。

3. SOD: 选择了五个数据集：DUTS [28], ECSSD[29], OMRON [30], HKU-IS [31], 和 PASCAL-S[32]。遵循了[33]中的训练协议，使用DUTS训练集进行训练，并在剩余的数据集上进行测试。

Inplementation Details.
  提出的解决方案在PyTorch中实现，所有训练都是在4块NVIDIA A40 GPU上进行的。训练了模型的两种版本：TS-SAM_B和TS-SAM_H。
TS-SAM_B使用SAM图像编码器的ViT-B版本，包括14层CSA和13层MRM。TS-SAM_H使用SAM图像编码器的ViT-H版本，包括34层CSA和13层MRM。所有实验都使用Adam优化算法，初始学习率为0.0008，并采用余弦衰减策略。总批大小为8。除了Resize之外，没有使用任何数据增强方法。对于COD和SOD任务，使用二值交叉熵（BCE）损失和IOU损失。训练进行了80轮。对于阴影检测任务，使用平衡二项交叉熵（BBCE）损失。训练进行了100轮。特别是，在阴影检测任务实验中，引入了图像的高频成分作为输入，遵循 [3]中的相同设置。

B. Results

  伪装目标检测COD. 表格I展示了TS-SAM与SAM、SAM-Adapter以及当前最先进的特定领域模型在四个常用伪装目标检测数据集上的结果。仅通过调整4.4%的参数，且没有为该任务设计特定的模块，所提模型在四个数据集上均取得了竞争性的性能。特别是在两个最大的数据集，COD10K和NC4K上，TS-SAM在大多数指标上取得了最佳结果。对于COD10K数据集，所提模型在$S_{alpha}$和$F_{\beta }^w$ 两个指标上分别超越了第二好的方法SARNet3.3%和4.4%，并且在M AE上降低了19.0%。在NC4K数据集上，所提模型在$S_{alpha}$、$F_{\beta }^w$和MAE三个指标上达到了当前最先进的性能，并且在$S_{alpha}$指标上超越了SARNet1.8%。这些结果展示了所提模型的强大泛化能力。

  此外，对于SAM-Adapter [3]，通过引入适配器，它在伪装目标检测任务上相对于原始的SAM取得了性能提升。而对于TS-SAM，通过将卷积侧适配器串联，将SAM的特征图融入到卷积侧适配器的训练中，并使用多尺度细化模块，它在SAM-Adapter[3]报告的三个数据集上所有指标上取得了显著的提升。

  图1展示了复杂样本的一些定性结果，展示了复杂性的各种方面。例如，存在由极其相似的环境和遮挡引起的干扰（第1、3和5列），小目标（第4列）和复杂纹理特征（第2列）。这些结果展示了提出的TS-SAM的优越性，包括对整体特征的良好掌握，对细节特征的深入探索，以及对干扰信息的有效区分。

  阴影检测。ISTD数据集上的结果见表II。观察到SAM-Adapter通过引入轻量级的适配器，在阴影检测任务中显著提高了SAM的性能。与SAM-Adapter相比，TS-SAM有显著的性能提升。而且，没有为阴影检测任务特别设计，TS-SAM达到了与最佳特定领域模型相近的性能。

  显著目标检测。表III展示了TS-SAM与SAM、SSOM[4]以及当前最先进的特定领域模型在五个常用显著目标检测数据集上的结果。与SSOM类似的方法SSOM相比，TS-SAM显示了显著的优势，在五个数据集上的MAE指标上显著超越了SSOM。这表明提出的微调方法优于SSOM提出的基于AdaLora [40]的微调方案。与高级特定领域方法相比，TS-SAM仅通过微调少量参数就达到了竞争性性能。在$S_{alpha}$、$E_{\varphi }$和MAE指标上，TS-SAM在ECSSD、OMRON和PASCAL-S数据集上达到了当前最先进的性能，并且在DUTS和HKU-IS数据集上表现良好。在$F^{w}_{\beta}$指标上，TS-SAM的表现相对较弱，这将是未来改进方向。

C. Ablation Study

 提出的模块的有效性在四个伪装目标检测数据集上进行了研究：CHAMELEON, CAMO, COD10K和 NC4K。TS-SAM 使用的模块在 SAM 的基线模型上单独实现，结果见表 IV。从结果来看，当只有解码器进行微调时，SAM 的基线模型表现不佳。首先，验证了独立工作的模块的有效性，包括只引入 CSA 模块和同时引入 MRM 和 FFD 模块的设置。发现与基线模型相比，这两种设置都带来了显著的性能提升，同时只增加了少量可训练参数（分别为 5.49M和 2.35M）。结果表明，提出的模块能够有效地从 SAM 图像编码器中提取视觉信息，并将其适应下游任务。其次，结合所有三个提议模块的模型在四个数据集上的所有指标上取得了最佳性能。结果表明，提出的模块可以互补地工作。专为分割任务设计的MRM 和 FFD，可以帮助 CSA 实现更好的性能。

五、总结

 文中专注于对大型视觉模型SAM下游任务的高效微调。为了充分利用SAM在大规模数据集上预训练的优势，引入了TS-SAM。使用轻量级的卷积侧适配器（CSA），首次将侧网络的概念应用于SAM的微调。此外，根据分割任务的特性，设计了**多尺度精炼模块（MRM）和特征融合解码器（FFD）**来从高分辨率图像中提取详细特征。在三个下游任务上的实验表明，所提模型超过了现有对SAM的高效微调方法，并且可以与专门为每个任务设计的SOTA领域特定模型相媲美。