【RS-Mamba】用于大型遥感图像密集预测的 RS-Mamba

摘要

• 上下文建模是遥感图像密集预测任务的关键。如今，超高分辨率（VHR）遥感图像的不断增长对有效建模提出了挑战。虽然基于transformer的模型具有全局建模能力，但由于其二次复杂度，它们在应用于大型VHR图像时遇到了计算挑战。传统的做法是将较大的图像裁剪成较小的patch，这样会导致上下文信息的显著丢失。
• 为此，提出RS-Mamba用于大型VHR遥感图像的密集预测任务。RS-Mamba是专门设计用于捕捉具有线性复杂度的遥感图像的全局上下文，促进对大型VHR图像的有效处理。RS-Mamba集成了一个全方位选择扫描模块，对多个方向图像的上下文进行全局建模，从不同方向捕捉大的空间特征。
• 研究证明了RSM在高分辨率遥感任务(语义分割和变化检测)中的高效性和优越性。利用线性复杂性和全局建模能力，RSM在大型遥感图像上实现了比基于transformer模型更好的效率和精度。还证明了该模型在密集的预测任务中，在更大的图像尺寸下通常表现得更好。
• 论文链接：https://arxiv.org/abs/2404.02668
• 代码链接：GitHub - walking-shadow/Official_Remote_Sensing_Mamba: Official code of Remote Sensing Mamba

方法

RS-Mamba总体架构如下图所示：

用于语义分割任务的遥感Mamba（RSM-SS）和用于变化检测任务的遥感Mamba（RSM-CD）

• RSM-SS架构
  • 使用U-Net编解码器框架，通过Patch Embedding将输入的超高分辨率遥感图像转换为一系列图像块，然后通过编码器提取特征，通过解码器上采样以产生语义分割结果。
  • 编码器包括五个阶段，每个阶段报考几个OSS块，阶段1负责从输入的VHR遥感图像中提取特征，阶段2-5逐步对编码器特征进行降采样，并将每个阶段的特征通道数量增加一倍。
  • 解码器由四个解码块组成，其中特征被上采样，然后通过跳跃连接和卷积，沿通道维度与编码器特征进行连接。该过程将解码器特征的语义信息和编码器特征的空间信息相结合，促进了从全局和局部两个角度进行语义分割。
• RSM-CD架构
  • 采用FC-Siam-Conc（Fully Convolutional Siamese Networks with Concatenation）的孪生网络架构，对双时相VHR遥感图像进行特征提取，然后将提取的特征进行融合和上采样，以获取变化检测结果。
  • 与RSM-SS类似，RSM-CD中的共享权重编码器由五个阶段组成，每个阶段有几个OSS块，并解码器由四个解码块组成。通过共享权重编码器提取特征后，将相同大小的双时态特征沿着通道维度进行连接并卷积。这种融合方式捕获了VHR遥感图像的两个时间段的信息，能够有效地分割变化的物体。将融合后的特征在编码器中进行上采样，通过跳跃连接和卷积与相同大小的融合特征进行连接，从而保持丰富的语义和空间信息。
• Omnidirectional State Space Block（OSS）
  • 如图2所示。OSS块的核心是面向扫描模块（OSSM），它是在图像中跨多个方向的全局上下文建模的核心。OSSM有选择地向不同的方向扫描输入图像，捕捉复杂的空间关系，并提供对上下文的全面理解。
  • OSSM将输入的tokens沿横向、纵向、斜向、反斜向及其反方向展平成8组序列，然后这些序列独立地通过SSM块进行选择性扫描。所有方向的扫描结果累加在一起形成输出tokens，这些输出tokens综合了8个方向上的特征，使模型能够全方位地捕获和建模遥感图像的大尺度空间特征。
  • 
  • 与Vim和VMamba对比，OSSM增加了斜向和反斜向扫描。
  • 

实验

Ablation Study

Overall Comparison

Semantic Segmentation Task:

Change Detection Task:

Impact of Image Size and Spatial Resolution

Handling Large Remote Sensing Images