Mamba入局遥感图像分割 | Samba: 首个基于SSM的遥感高分图像语义分割框架

标题：《Samba: Semantic Segmentation of Remotely Sensed Images with State Space Model》 论文：https://arxiv.org/abs/2404.01705 源码：https://github.com/zhuqinfeng1999/Samba

1、导读

本文将为大家介绍来自西交利浦大学、澳大利亚科学与工业研究院矿物研究所以及利物浦大学联合提出的最新工作Samba，这是一种基于状态空间模型的遥感图像语义分割方法。该工作提出了高分辨率遥感图像的语义分割框架 Samba，性能指标均超越了CNN-based和ViT-based方法，证明了Mamba在遥感高分图像中作为新一代backbone的潜力，为遥感语义分割任务提供了 mamba-based 方法的表现基准。值得一提的是，Samba是第一个将状态空间模型（SSM）运用到遥感图像语义分割任务中的工作，代码已开源，欢迎关注。

2、背景

高分辨率遥感图像对常用的语义分割方法提出了挑战，如卷积神经网络（CNN）和Vision Transformer（ViT）。CNN-based方法难以处理这种高分辨率图像，因为它们的感受野受限，而ViT面临处理长序列的挑战。

受到状态空间模型（SSM），也叫Mamba，能够有效捕获全局语义信息的启发，我们提出了一个针对高分辨率遥感图像的语义分割框架，命名为Samba。

Samba利用编码器-解码器架构，其中Samba块作为编码器用于多级语义信息的有效提取，而UperNet作为解码器。

我们在LoveDA数据集上对Samba进行了评估，并将其性能与表现最佳的CNN和ViT方法进行了对比。结果显示，Samba在LoveDA上取得了杰出的表现。这表明，所提出的Samba框架是SSM在遥感图像语义分割中的有效应用，为这一特定应用中基于Mamba技术的设定了新的性能基准。

3、动机

遥感图像语义分割是遥感应用中的一项关键任务，深度学习方法被广泛应用。在这些方法中，常用的深度学习技术是卷积神经网络（CNN）。通过执行滑覆盖图像数据的卷积操作，CNN能够有效地从图像的浅层到深层提取语义特征，成为许多图像处理任务的基石。

然而，如图1（a）所示，CNN中有限的感受野在处理高分辨率图像时呈现出挑战。

尽管存在缓解该问题的解决方案，但它们都有不可避免的缺陷。例如：

1. 缩放图像能够适应网络的感受野，但通常会导致图像分辨率丢失，从而影响模型性能。
2. 空洞卷积可以扩大感受野，但由于粗糙的特征子采样，可能导致特征信息损失。
3. 另一种方法是通过残差连接多个CNN，将高级语义与低级信息融合，从而增强模型识别不同尺度的能力。然而，该方法显著增加了网络计算复杂度。

此外，ViT是另一种语义分割广泛采用的深度学习技术。凭借其全局注意力机制，ViT克服了感受野的限制，能够灵活适应不同分辨率的输入，如图1（b）所示。

ViT 在传统图像任务中展现了卓越的性能，例如在IamgeNet数据集中的分类任务。然而，在采用ViT进行遥感图像的语义分割时，仍然存在挑战。由于需要对每个图像块之间的自注意力进行计算，随着分辨率的增加，计算复杂度呈指数级增长。不仅如此，由于缺乏归纳偏置，ViT需要大量的训练数据，但遥感领域的图像数据通常是稀缺的，同时对数据增强的要求也很高。

近期，名为Mamba的新方法被提出了，该方法利用状态空间模型（SSM）以低计算复杂度捕获全局语义信息。与Transformer不同，Mamba展示了线性复杂度，在处理长序列问题中体现了明显优势。在视觉任务中，探索将Mamba替换多头自注意力的效果是一个有趣的研究方向。

因此，受Mamba启发，我们提出了Samba，一个为遥感高分辨率图像设计的语义分割框架。Samba块被设计用于高效的图像特征提取。在该框架中，Mamba取代了ViT中的多头自注意力来捕获图像数据信息，并与多个多层感知器（MLP）结合，创建了一个Samba块。所提出的语义分割框架采用编码器-解码器架构，使用Samba块作为编码器和UperNet作为解码器，以有效提取多级语义信息。

我们的方法的表现性能使用LoveDA数据集进行了评估。与表现最佳的CNN进而ViT方法相比，在未加载预训练参数的情况下，Samba在LoveDA上展现了卓越的性能表现。这表明，Samba是状态空间模型在遥感图像语义分割中的有效应用，为该领域中基于Mamba的方法设立了新的基线。

本研究的主要贡献可以总结如下：

1. 我们提出了Samba架构，首次将Mamba架构引入遥感图像的语义分割任务中。
2. 我们进行了与表现最佳网络的对比实验，展示了Mamba架构作为遥感图像语义分割骨干网络的巨大潜力。
3. 我们为基于Mamba的遥感图像分割建立了新的性能基准，并提供了见解以及提出了潜在的未来方向。

4、方法

图2 展示了Samba的编码器架构：

该架构中，Samba块在四个阶段中对图像进行逐层的下采样。ViT编码器使用多头自注意力来捕捉不同表示子空间内的信息，随后通过残差连接和层归一化（LN）来缓解梯度消失问题。接着，采用一个前馈网络（FFN），由一个MLP和LN组成，以引入非线性变换并整合由多头自注意力精炼的复杂信息。

受到ViT编码器稳健的架构的启发，Samba块采用了类似的架构，多头自注意力被替换成了Mamba块。Mamba块被用于从高分辨率图像序列中提取特征，避免了在计算中的二次复杂性。我们的方法采用了Mamba块和MLP的结合，以增强模型的表示能力并加强其对复杂数据的学习能力。

Mamba的核心状态空间模型可以用演化参数A，投影参数B和C的线性微分方程表示：

状态空间模型的离散化：

更多细节请阅读原始论文。

5、实验

LoveDA数据集被采用对我们提出的框架进行验证，该数据集包含具有0.3米空间分辨率的遥感图像。它包括2522幅训练图像、1669幅验证图像和1796幅测试图像，其中验证集用于性能评估。平均交并比（mIoU）被用来评估分割的准确性。

在这项研究中，Samba与几种表现最佳的方法进行了比较。包括CNN-based的方法，如ConvNeXt、ResNet、Deeplab V3+和PSPNet，以及ViT-based的方法，如Swin-T。为了确保公平比较，测试的方法都未加载预训练参数。这些方法的优化器和学习率策略设置遵循广泛采用的最佳配置。我们使用随机调整大小、随机裁剪、随机翻转和光度畸变来增广训练数据。具体的训练设置在表1中总结。所有实验均使用两块NVIDIA RTX 3090和两块4090进行。

具体来说，，当使用UperNet作为解码器时，Samba在mIoU指标上以2.24%的优势超越了性能最佳的ViT-based模型Swin-T，并以6.51%的优势超越了了性能最佳的CNN-based的模型ConvNeXt。当与UperNet解码器组合时，Samba在每个patch的flops和参数上都比Swin-T、ResNet50和ConvNeXt要低。

此外，如图3所示，我们将Samba在LoveDA数据集上实现的语义分割结果与对比模型的结果进行了可视化。这些可视化进一步展示了Samba在高分辨率遥感图像语义分割中的性能优势。

6、总结

本文介绍了Samba，一种基于Mamba的高分辨率遥感图像语义分割框架，标志着Mamba在该领域的首次应用。通过在LoveDA数据集上性能的评估，Samba超越了最先进的CNN-based和ViT-based的方法，设定了新的性能基准，并展示了Mamba架构在高分辨率遥感影像语义分割中的有效性和潜力。

参考：CVHub 作者-派派星