【2024| Mamba 】遥感图像语义分割——RS^3 Mamba！！论文解读

【遥感图像语义分割SSRS】$R{S}^{3}Mamba$: Visual State Space Model for Remote Sensing Image Semantic Segmentation

$R{S}^{3}Mamba$: Visual State Space Model for Remote Sensing Image Semantic Segmentation

---

论文链接：https://ieeexplore.ieee.org/abstract/document/10556777
代码链接：https://github.com/sstary/SSRS

---

摘要

• 遥感图像语义分割是地学研究的一项基本任务。然而，卷积神经网络(cnn)和transformer有一些显著的缺点。前者受到长程（long-range）建模能力不足的限制，而后者受到计算复杂性的阻碍。最近，出现了一种以Mamba为代表的新的视觉状态空间(visual state space,VSS)模型，能够以线性可计算性建模长程关系。提出了一种专为遥感任务设计的双分支网络$R{S}^{3}Mamba$(remote image semantic segmentation Mamba)。$R{S}^{3}Mamba$使用VSS块来构建辅助分支，为基于卷积的主分支提供额外的全局信息。此外，考虑到两个分支的不同特点，本文引入了一个协同补全模块(collaborative completion module,CCM)，使用一种新的自适应机制来细化和融合来自双编码器的特征。通过在两个广泛使用的数据集上进行实验，$R{S}^{3}Mamba$在ISPRS Vaihingen和LoveDA Urban上的mIoU分别达到0.66%和1.70%，优于目前最好的方法，证明了其有效性和潜力。
  
  $R{S}^{3}Mamba$总体架构。它由VSS辅助编码器、残差主编码器(包括用于跨分支语义融合的CCM)和解码器组成。

---

论文发表单位：
1.香港中文大学 未来情报网络研究所
2.武汉科技大学 信息科学与工程学院
3.香港中文大学 理工科学院

一、引言

• 现代地球科学研究在很大程度上依赖于卫星或航空设备收集的大量遥感数据。这些数据捕捉了地球表面物体的光谱特征，并提供了自然和人造结构的准确视觉表示。语义分割方法旨在将遥感图像中的每个像素划分为不同的类别，从而帮助研究人员全面探索地表条件。这种自动化分析和解释方法支持各种下游任务和应用，包括土地覆盖制图和环境监测${}^{[1]}$。
• 近几年，基于深度学习的方法作为一种数据驱动的自动化技术，已经彻底改变了语义分割${}^{[2]}$。目前，遥感领域的主流模型主要分为两类：CNN和Transformer。前者可以通过堆叠卷积操作提取图像特征，后者可以基于自注意力机制${}^{[3,4]}$对长程依赖关系进行建模。尽管具有这些优点，但这些模型在应用于遥感图像时存在局限性。与自然图像相比，遥感图像通常表示包含多个目标的复杂场景，且目标尺度变化显著。因此，CNN受到其局部感受野的限制，使其难以掌握和学习复杂的表示。另一方面，尽管Transformer具有学习长程依赖关系的能力，但其高计算复杂度对模型效率和内存占用提出了重大挑战。
• 最近，建立在状态空间模型(SSM)上的Mamba${}^{[5,6]}$被提出。为了平衡模型的效率和效果，引入选择机制来控制信息沿着特定序列维度的传播或交互。通过让影响序列交互的参数依赖于输入，该模型可以建立长距离依赖关系，同时保持线性的计算复杂度${}^{[5]}$。在遥感领域，He${}^{[7]}$等人提出了Pan-Mamba，包括信道交换Mamba和跨模态Mamba，用于Pan-Mamba锐化。Chen等人${}^{[8]}$提出了RS Mamba，其中包括一个用于大规模图像解释的多路径视觉状态空间(VSS)块。这两种方法直接用VSS块替换现有网络。然而，由于大多数现有的基于VSS的模型是从头开始训练的，与完全预训练的Cnn和Transformer相比，实现良好的性能提出了更大的挑战。
• 为了将VSS模块引入到遥感图像语义分割中，应对上述挑战，本文提出了一种辅助分支策略，利用VSS模块提供额外的全局信息，辅助基于卷积的主分支进行特征提取。为解决全局和局部语义之间的差异，引入了一个协同补全模块(CCM)，以促进跨分支语义融合。本文的研究贡献如下:

（1）本文提出$R{S}^{3}Mamba$，标志着基于VSS模型在遥感图像语义分割中的开创性探索。这为构建具有建模长距离关系能力，同时保持遥感任务线性复杂度的模型提供了有价值的见解。

（2）本文提出一种新颖的CCM方法来融合跨分支特征。CCM通过自适应设计减小全局和局部语义之间的差异，并将跨分支特征投影到同一特征子空间中，实现特征的全融合。

二、预备工作

$SS2D$数据流。它将输入扩展到四个方向，通过$S6$逐个扫描，然后合并它们

• SSM是一个线性时不变系统，通过隐藏状态$h(t)\in R^N$，将输入$x(t)\in R^L$投影到输出$y(t)\in R^L$。在数学上，它可以用线性常微分方程(ODEs)表示为：
  
  其中，$A\in R^{N\times N}$表示状态转移矩阵，$B\in R^{N\times 1}$, $C\in R^{1\times N}$表示状态大小为$N$的投影矩阵。
• 在深度学习方法中，ODEs必须离散化。选择性扫描空间状态序列模型(selective scan space state sequential
  model, $S6$)${}^{[5]}$是连续系统的离散版本。$△$是一个时间尺度参数，通过常用的零阶保持规则和一阶泰勒级数将连续参数$A、B$转换为离散参数$\overline{A}、\overline{B}$：

• (2)离散化后的ODEs可重写为：

• 在S6中，矩阵$B、C$和$△$是从输入$x_k$${}^{[5]}$推导出来的。
  ${}^{[6]}$将SSM引入到视觉任务中，并给出了图1所示的二维选择性扫描(SS2D)。它在四个方向上扩展图像块，以创建四个单独的序列。这些序列然后通过SSM单独处理。最后，将得到的特征进行组合，生成完整的2-D特征图。给定输入特征映射$x$, SS2D的输出特征映射$\overline{x}$可以表示为：

其中$v\in 1,2,3,4$表示四个不同的扫描方向。此外，expand$(\cdot )$和merge$(\cdot )$分别表示扫描展开和扫描合并操作。

• (5)中描述的S6函数作为VSS块的核心算子，促进了一维数组中的每个元素与任何先前扫描的样本之间通过压缩隐藏状态的相互作用。

三、方法详解

$R{S}^{3}Mamba$由三部分组成，即VSS辅助编码器、残差主编码器包括用于跨分支语义融合的CCM，以及解码器如图2所示。具体而言，$R{S}^{3}Mamba$模型通过主、辅助分支提取图像特征，每个分支由4个相应的块组成。将辅助分支生成的特征输入到主分支内对应尺度的CCM中进行特征融合。在4个尺度上进行特征提取和特征融合后，得到多尺度特征，通过跳跃连接输入到解码器生成预测图。本研究采用UNetformer${}^{[3]}$中的解码器。

(a)以SS2D为核心计算单元的VSS模块详细架构。(b) CCM的详细结构。两个分支用于融合跨分支的特征映射。

Auxiliary Encoder

• 辅助编码器基于模拟长程依赖${}^{[9]}$的VSS块构造。具体结构如图3(a)所示，其中SS2D是VSS模块的核心计算单元。如图2所示，辅助编码器由四个连续的阶段组成，每个阶段由patch算子和VSS块组成。第一阶段包含一个块嵌入层和VSS块，而接下来的三个阶段分别包含一个块合并层和VSS块。我们用$X\in R^{H\times W\times Z}$表示输入图像，$Z$是图像通道的数量。$H$和$W$分别表示图像的高度和宽度。patch embedding层首先将$X$分割成互不重叠的patch，然后使用VSS块提取特征$F_a^1$。之后，编码器通过执行另外三个阶段的类似操作来生成$F_a^{2-4}$。值得注意的是，由于辅助分支的特征流不受主分支的影响，因此整个辅助分支可以独立应用。

Main Encoder and CCM

采用ResNet18作为主编码器学习局部表示；如图2所示，它包含4个残差块和4个CCMs。四个残差块进行卷积操作，生成多尺度特征$F_m^{1-4}$。与辅助编码器相比，主编码器利用现有的预训练模型有效地从遥感图像中提取特征。在此基础上，通过CCM将辅助分支提取的不同尺度的特征融合到主分支中，弥补了主分支在提取全局信息方面的局限性。
所提CCM的详细结构如图3(b)所示。所提出的CCM用于融合$i\in 1,2,3,4$的跨分支特征图$F_m^i$和$F_a^i$。具体来说，CCM由两个平行的分支组成，即全局分支和本地分支。前者用于从主支路增强$F_a^i$，后者用于从辅助支路加工$F_a^i$。考虑到主分支的特征是通过具有局部特性的卷积操作获得的，使用基于窗口的多头自注意力对长程依赖进行建模。需要强调的是，该机制还保持了线性复杂度。另一方面，考虑到辅助分支的特征是由具有长程特性的VSS获得的，因此使用卷积学习局部细节。通过对主要局部特征进行全局补全，对辅助全局特征进行局部补全，将它们投影到相同的特征子空间中，以实现更好的特征融合。考虑到双分支特征的完备性，将该自适应融合模块命名为CCM。
在主分支的每个阶段，CCM都在相应的尺度上进行特征融合。然后，将得到的融合特征(记为$F_r^i$)作为跳跃连接送入解码器。解码器在生成最终预测图之前恢复抽象特征。

四、实验和讨论

数据集

1. ISPRS Vaihingen: ISPRS Vaihingen数据集由16张真实正射影像组成，每张都具有非常精细的分辨率，平均大小为2500 × 2000像素。这些正射影像由三个通道组成:近红外、红、绿(NIRRG)，地面采样距离为9 cm。该数据集包含5个前景类，即不透水的表面、建筑物、低植被、树、汽车和一个背景类。将真实正射影像裁剪为256 × 256，然后分为960个训练块和320个测试块。
2. LoveDA Urban: LoveDA数据集提供了一个全面的场景集合，包括城市和农村。在本研究中，由于地物分布的多样性，我们选择了LoveDA城市场景。LoveDA城市场景由1833幅1024 × 1024像素的精细光学遥感图像组成。每幅图像提供红“绿”蓝(RGB)三个通道，地面采样距离为30 cm。该数据集包含7个土地覆盖类别，包括背景、建筑、道路、水、荒地、森林和农业。将1833幅图像分为1156幅图像的训练集和677幅图像的测试集。

• 这两个数据集在采样分辨率、地物类别和标签精度方面进行了对比。通过在两个数据集上进行实验，旨在证明$R{S}^{3}Mamba$的有效性和优越性。

实验设置

• 采用基于ResNet18的UNetformer${}^{[3]}$作为基线模型。所提出的$R{S}^{3}Mamba$与几种最先进的监督方法进行了基准测试，包括UNetformer${}^{[3]}$， CMTFNet${}^{[10]}$，TransUNet${}^{[11]}$和SSNet${}^{[12]}$。根据采用的backbone，前两种方法基于CNN，后两种方法基于Transformer。
• 实验使用PyTorch在单个配备24-GB RAM的NVIDIA GeForce RTX 4090 GPU上进行。采用随机梯度下降(SGD)算法对所有模型进行优化训练。学习率为0.01，动量为0.9，衰减系数为0.0005，批量大小为10。总epoch被设置为50，每个epoch一个测试。训练和测试的输入大小都是256 × 256。在定量评估方面，使用了两个广泛使用的指标:平均F1-score (mF1)和平均交并比(mIoU)。加粗的值表示表中表现最好的指标。

性能对比

Performance Comparison on the Vaihingen Dataset

• 如表1所示，与基线UNetformer相比，$R{S}^{3}Mamba$在mF1和mIoU方面都有很大的提升。这证实了基于VSS块的双分支架构可以有效增强特征提取。实验结果表明，$R{S}^{3}Mamba$模型在不透水表面、低植被、树木和汽车4个类别上表现优于现有的主流模型。与CMTFNet相比，$R{S}^{3}Mamba$在不透水地表类别的F1和IoU指标上分别提高了0.30%和0.53%。与基线UNetformer相比，低植被的IoU提高了0.51%，建筑物的IoU提高了0.44%。在综合性能方面，$R{S}^{3}Mamba$取得了90.34%的mF1分数和82.78%的mIoU，分别比基线UNetformer的相应性能提升了0.49%和0.81%。这些增加可以归因于辅助分支提供的全局语义信息和自适应跨分支语义融合。图4说明了所考虑的所有五种方法所得结果的可视化示例。可以看出$R{S}^{3}Mamba$能够更准确地分割出边界更平滑、噪声更少的地物。
  

Performance Comparison on the LoveDA Urban

• 尽管两个数据集在样本分辨率和地物类别上存在差异，但在LoveDA Urban数据集上的实验与ISPRS Vaihingen数据集的实验结果相似，如表II所示。特别是，该方法显著提高了农业类别的准确性，这对现有方法来说是典型的挑战。在综合性能方面，$R{S}^{3}Mamba$在mF1得分和mIoU上分别比基线提升了1.52%和1.81%。图5展示了LoveDA Urban的可视化示例，可以清楚地表明对地物的更彻底的识别。这证实了利用全局辅助信息进行遥感图像表示学习以及跨全局和局部语义特征融合的有效性和必要性。
  

消融实验

• 为了测试$R{S}^{3}Mamba$中提出的VSS辅助分支和CCM模块的有效性，进行了4个消融实验，如表III所示。对表III中的前两行数据的检查表明，与经典的ResNet相比，现有的基于VSS的backbone网络的性能不太令人满意，原因是预训练不足。第三行的结果进一步验证了双分支结构的有效性。作为辅助分支，VSS编码器可以提供额外的信息来辅助特征提取和最终的语义恢复。第三行和第四行实验结果验证了CCM模块的有效性，该模块根据双分支编码器的结构特点进一步有效挖掘特征。表3证实，$R{S}^{3}Mamba$是探索基于VSS的遥感任务方法的有效框架。
  

模型复杂度分析

• $R{S}^{3}Mamba$的计算复杂度使用三个指标进行评估：浮点操作计数(FLOPs)、模型参数和内存占用。FLOPs访问模型的复杂性，而模型参数和内存占用分别评估网络的规模和内存需求。理想的模型应该表现出较小的浮点数、模型参数和内存占用值。
• 表4给出了本研究中考虑的所有语义分割模型的复杂度分析结果。对表IV的检查表明，与基线UNetformer相比，$R{S}^{3}Mamba$引入了额外的模型复杂性。这是由于合并了一个额外的辅助分支，并指定了特定的跨分支特征融合模块。然而，值得注意的是，与基于Transformer的方法相比，所提出方法在计算复杂度和模型规模上有显著降低。这些结果突出了所提出方法的有效性，作为将Mamba引入遥感任务的实用方法。在Mamba发展的早期阶段，该方法可以为Mamba在该领域的未来发展提供有价值的见解。
  

五、总结

针对遥感图像语义分割任务，提出了一种基于VSS的$R{S}^{3}Mamba$模型。特别地，开发了一个基于VSS块的辅助分支，以最小的线性计算复杂度提供额外的全局感知信息。此外，设计了一个新的CCM模块，利用主分支和辅助分支产生的不同特征进行跨分支特征融合。与现有的直接用完整的VSS模型替换CNN和Transformer的方法相比，$R{S}^{3}Mamba$为将Mamba集成到遥感图像语义分割模型中提供了一种独特而有效的方法。在两个成熟的遥感数据集上进行的广泛实验表明，$R{S}^{3}Mamba$优于其他最新的基于CNN和Transformer的语义分割模型，且计算复杂度合理。我们希望这项研究能激励人们继续研究$Mamba$在遥感领域的应用。