Mamba入局遥感！RS-Mamba：首次使用SSM进行大遥感图像的语义分割与变化检测

©作者 | 赵思杰

单位 | 南京大学硕士生

研究方向 | 语义分割

Mamba 在大语言模型大放异彩，以其线性复杂度和媲美 transformer 的表现，被认为是 transformer 的有力替代。近期工作 Vim 和 VMamba 将 Mamba 引入到视觉图像领域，引爆了视觉领域的众多领域，涌现出大量使用 Mamba 进行视觉任务的研究。

本文首次将 Mamba 引入到遥感，构建了 RS-Mamba 来进行超高分辨率遥感图像的密集预测任务，利用它的线性复杂度和全局建模能力来处理大遥感图像。

之前的遥感模型主要可以分为基于 CNN 和基于 transformer。基于 CNN 的模型由于局部卷积操作，无法对遥感图像进行全局建模，基于 transformer 的模型由于二次方复杂度无法处理大的超高分辨率遥感图像，将大的图像裁剪为小的图像块会损失大量的上下文信息。

近期的工作 Vim 和 VMamba 虽然将 Mamba 引入了视觉图像领域，但是它们只在图像的横向或者纵向进行选择性扫描，适合主要空间特征分布在横向或者纵向的自然图像，但是不适合空间特征分布在任意方向的遥感图像。

针对上述问题，RS-Mamba 创新性的提出了全向选择性扫描模块，在多个方向对遥感图像进行选择性扫描，从而能够提取出多个方向的大尺度空间特征。同时由于 RS-Mamba 只具有线性复杂度，它能够处理基于 transformer 模型无法处理的大遥感图像，并具有全局建模能力。

在多种地物的语义分割和变化检测任务上的实验表明，使用简单的模型架构和训练方式，RS-Mamba 就能够达到 SOTA。

▲ RS-Mamba总体结构

代码链接：

https://github.com/walking-shadow/Official_Remote_Sensing_Mamba

论文链接：

https://arxiv.org/abs/2404.02668

代码已开源，觉得有帮助的话希望能在 github 给个 star。

欢迎在本文的基础上继续探究基于 SSM 的方法在遥感密集预测任务中的潜力，RSM 目前所使用的都是最简单的架构，还有很大的潜力可以挖掘。

Mamba 在各个领域的火热之风很快就会吹到遥感领域，它在遥感领域的潜力想必也会迎来新一轮的研究热潮。

引言

近年来，遥感图像的空间分辨率正在变得越来越高，超高分辨率遥感图像也广泛应用在大量的遥感领域中。

超高分辨率遥感图像具有多个方向的大空间尺度的空间特征，它们对于语义分割、变化检测等密集预测任务来说都至关重要。

在超高分辨率遥感图像中，由于图像的空间分辨率非常高，一方面，单个物体内部存在形状、边缘等的丰富的空间特征，多个物体之间也存在空间排列等丰富的空间特征，这些空间特征往往具有很大的空间尺度；另一方面，由于遥感图像是相机向下俯拍得到的，水平方向上相机可以以任意角度获取遥感图像，这表明遥感图像的空间特征可以存在于任意方向。

因此，能够对超高分辨率遥感图像进行全局建模，提取多个方向的大尺度空间特征，对于超高分辨率遥感中的密集预测任务来说至关重要。

近年来，基于 transformer 的深度学习模型被广泛应用于超高分辨率遥感任务中。由于 transformer 能够通过自注意力捕获超高分辨率遥感图像的全局空间信息，并对它们的空间依赖进行良好建模，基于 transformer 的模型取得了很不错的效果。

然而，由于 transformer 具有二次方复杂度，在训练和推理基于 transformer 的模型时，只能将大的超高分辨率遥感图像，裁剪成较小的遥感图像块，然后再送入到基于 transformer 的模型中进行计算。

由于超高分辨率遥感图像的物体存在很大的空间跨度，裁剪后的遥感图像块往往只包含单个物体的一部分，只具有更少的上下文信息，丢失了单个物体内部的空间特征和多个物体之间的空间依赖信息，不利于处理超高分辨率遥感任务，如下图所示。

▲ 把大遥感图像裁剪为小的图像块。小图像块包含的上下文信息非常有限，丢失了很多重要的空间特征

最近的工作 Vim 和 VMamba 使用 SSM 实现了线性复杂度和全局感受野，在自然图像上完成图像分类、图像分割等任务。

为了处理图像数据无方向性的特点，Vim 在图像的横向方向上使用 SSM 进行前向和后向选择性扫描，VMamba 在横向和纵向方向上使用 SSM 进行前向和后向选择性扫描，从而使得图像中的每个部分都能和其他部分建立联系。

VMamba 中的有效感受野可视化结果表明，它具有全局的有效感受野，并且在横向和纵向方向上具有更强的有效感受野，这表明 SSM 的选择性扫描方向能够显著影响图像在特定方向的有效感受野。

然而，Vim 和 VMamba 不适用于超高分辨率遥感图像。自然图像的空间特征要符合一定的物理规律，图像中的物体往往是横平竖直的，因此自然图像不能随意旋转，它的主要空间特征分布在横向和纵向两个方向。

遥感图像由于是卫星向下俯视拍摄获取的，因此遥感图像可以任意旋转，它的主要空间特征可以分布在任意方向。由于超高分辨率遥感图像中的物体具有大的空间尺度，单个物体的空间特征和多个物体之间的空间依赖都存在多种多样的方向，因此超高分辨率遥感图像存在多个方向的大尺度空间特征。

由于 SSM 的选择性扫描方向能够显著影响图像在特定方向上的有效感受野，Vim 只在横向上选择性扫描图像，VMamba 只在横向和纵向上选择性扫描图像，虽然它们在主要空间特征在横向和纵向的自然图像上能够取得不错的效果，但是它们都无法处理超高分辨率遥感图像中具有多个方向的大尺度空间特征

为了解决上述问题，我们首次将 SSM 引入超高分辨率遥感任务以实现全局感受野和线性复杂度，提出了 Remote Sensing Mamba（RSM）来处理超分辨率遥感图像。

RSM 不包含任何自注意力操作，但是具有全局感受野，能够对超高分辨率遥感图像的上下文进行有效建模。由于 RSM 具有线性复杂度，它能够处理包含了多个物体的大尺度的超高分辨率遥感图像，不会丢失单个物体的空间特征信息和多个物体之间的空间依赖信息，因此可以很好的处理超高分辨率遥感图像。

同时，我们提出了 Omnidirectional selective scan module（OSSM），来提取超高分辨率遥感图像中具有大空间尺度和多个方向的空间特征。OSSM 同时在横向、纵向、斜向和反斜向使用 SSM 对遥感图像进行前向和后向选择性扫描，从而能够增强遥感图像在多个方向上的全局有效感受野，在多个方向上提取全局的空间特征。

总的来说，我们的贡献如下：

1）我们提出了 Remote Sensing Mamba 来处理超高分辨率遥感任务。RSM 首次使用 SSM 来处理超高分辨率遥感图像，它能够处理包含整个物体的超高分辨率遥感图像，并建立起遥感图像的全局联系。

2）我们设计了一个 Omnidirectional selective scan module 提取超高分辨率遥感图像中具有大空间尺度和多个方向的空间特征。OSSM 通过在多个方向上使用 SSM 对遥感图像进行选择性扫描，能够增强遥感图像在多个方向上的全局联系。

3）我们证明了 RSM 在超高分辨率遥感任务中的高效性和优越性。在语义分割数据集 WHU-SS 和变化检测数据集 LEVIR-CD，和 WHU-CD 上的实验表明，RSM 在使用简单的模型架构和训练方法的情况下，在语义分割和变化检测任务上均能够达到 SOTA。

研究方法

RSM 在语义分割和变化检测任务上的对应模型分别为 RSM-SS 和 RSM-CD。RSM-SS 使用简单的 U-Net 架构，RSM-CD 使用简单的孪生网络架构，它们都是遥感密集预测任务中非常常用和简单的架构，不包含任何花哨的模块。

RSM-SS 和 RSM-CD 的 encoder 都由若干全向状态空间块（Omnidirectional state space block, OSS）构成，OSS 的内部结构为类 Mamba 结构，具有线性复杂度和全局建模能力。而其中的 OSSM（Omnidirectional selective scan module）在八个方向上对图像进行选择性扫描，从而能够提取多个方向的大尺度空间特征。

▲ RSM-SS 和 RSM-CD 的总体结构

RSM-SS 和 RSM-CD 的创新之处在于 OSSM 模块，OSSM 的扫描方向能够更好的适应遥感图像的空间特征分布在任意方向的特点。

Vim, VMamba 和 OSSM 的选择性扫描方向如下图所示，Vim 的选择性扫描方向是图像的横向，即从左到右扫描和从右到左扫描；VMamba 的选择性扫描方向是图像的横向和纵向，即从左到右和从右到左，加上从上到下和从下到上扫描。

Vim 和 VMamba 的这种扫描方式能够有效的处理主要空间特征在横向和纵向的自然图像，但是无法处理空间特征分布在任意方向的遥感图像。因此，OSSM 在横向和纵向的扫描基础上，增加了斜向和反斜向的扫描，从而能够在斜向和反斜向方向对遥感图像进行全局建模，提取多个方向的空间特征。】sdfdsfdfsd

▲ Vim, VMamba 和 OSSM的选择性扫描方向

OSSM 的具体结构如下。输入的 tokens 在横向、纵向、斜向、反斜向和它们的反向方向展平成 8 个图像 token 序列，并独立的输入到 SSM 块中进行选择性扫描，所有方向的扫描结果最后加在一起，得到输出的 tokens。输出 tokens 因为汇总了在 8 个方向上进行选择性扫描的特征，能够在多个方向上对遥感图像进行全局性建模，从而能够提取多个方向的大尺度空间特征。