0


Mamba~合集1

# PyramidMamba

从 CNN 到 Mamba|PyramidMamba 网络的语义分割优势,解决多尺度特征融合中的语义冗余问题

本文设计了一个即插即用的解码器,该解码器开发了一种密集空间金字塔池化(DSPP)来编码丰富的多尺度语义特征,以及一个金字塔融合曼巴(PFM)来减少多尺度特征融合中的语义冗余。全面的消融实验证明了所提出方法在增强多尺度特征表示以及实时语义分割巨大潜力方面的有效性和优越性。

在深度学习(DL)浪潮的推动下,基于CNN和Transformer的语义分割方法得到了广泛探索,这两种架构都揭示了多尺度特征表示对于加强地理目标的语义信息的重要性。然而,实际的多尺度特征融合通常伴随着由于金字塔特征中的同质语义内容而产生的语义冗余问题。

为了解决这个问题,作者提出了一个基于曼巴的分割网络,名为金字塔曼巴( PyramidMamba)。具体来说,作者设计了一个即插即用的解码器,该解码器开发了一种密集空间金字塔池化(DSPP)来编码丰富的多尺度语义特征,以及一个金字塔融合曼巴(PFM)来减少多尺度特征融合中的语义冗余。全面的消融实验证明了所提出方法在增强多尺度特征表示以及实时语义分割巨大潜力方面的有效性和优越性。此外,作者的金字塔曼巴在三个公开可用的数据集上取得了最先进的表现,即OpenEarthMap(70.8% mIoU)、ISPRS Vaihingen(84.8% mIoU)和Potsdam(88.0% mIoU)数据集。

代码将在https://github.com/WangLib91995/GeoSeg提供。

I Introduction

高分辨率遥感图像的语义分割对于地球观测(EO)的广泛应用越来越重要[1],如土地利用和覆盖(LULC)制图[2, 3, 4],环境监测[5, 6],以及城市可持续发展[7, 8]。在人工智能和传感器技术的推动下,深度学习(DL)[9]已经无缝地融入遥感领域,成为处理和分析大量遥感大数据[10, 11]的催化剂。与传统机器学习方法,如支持向量机(SVMs)和随机森林相比,基于DL的方法在自动和健壮的语义特征提取方面表现出优势,从而获得更高的准确性。

全卷积网络(FCN)[12]首次采用基于DL的端到端方式构建用于语义分割的新型卷积神经网络(CNN)。从那时起,基于FCN的语义分割方法逐渐成为主流。尽管FCN取得了重要的突破,但其单一和有限的感受野导致分割粗糙[13]。为了解决这个问题,一些研究应用了注意力机制来引入全局上下文信息。典型的方法包括非局部神经网络[14],双注意力网络[15]和Transformers[16, 17]。另一部分研究采用金字塔或多尺度特征融合方案以实现多感受野。最具代表性的方法包括金字塔场景解析网络(PSPNet)[18]和特征金字塔网络(FPN)[19]。然而,上述方法都有其不足之处。由于注意力机制的计算复杂性,前者需要大量计算资源来建模全局信息,效率较低。后者经常导致多尺度特征融合中的冗余问题,因为金字塔特征中存在大量同质语义信息。因此,如何有效地聚合多尺度语义特征仍然是一个挑战。

最近,一种基于选择性状态空间模型(SSM)[20]的新型架构,即Mamba,在计算机视觉和自然语言处理领域引起了广泛关注。与为序列建模应用低效的自注意力机制的Transformers不同,Mamba利用了选择性扫描机制,独特的硬件感知算法和并行扫描,在处理长序列方面具有高效率的优势。特别是,选择性扫描机制允许Mamba压缩同质特征并提取核心语义信息。因此,这个方案在解决多尺度特征融合中的冗余问题方面具有巨大潜力。

在本文中,作者构建了一个基于Mamba的网络,即PyramidMamba,用于遥感图像的语义分割。具体来说,作者开发了一个基于Mamba的解码器,包括一个密集空间金字塔池化(DSPP)模块和一个金字塔融合Mamba(PFM)模块。DSPP允许与标准空间金字塔池化模块相比更多的池化尺度,从而捕捉更多细粒度的多尺度上下文。PFM引入标准的Mamba块以聚合金字塔语义特征,减轻冗余问题并增强多尺度视觉表示。此外,即插即用的DSPP和PFM可以集成到深度神经网络中,用于高效有效的多尺度特征表示。本文的主要贡献可以总结如下:

  1. 作者重新思考了金字塔特征融合方案,并开发了一种新型的基于Mamba的分割网络(PyramidMamba)以改进多尺度特征表示。
  2. 作者设计了一个基于Mamba的解码器,应用密集空间池化以产生更细粒度的多尺度上下文,同时使用Mamba的选择性特征有效减少多尺度特征融合中的同质语义信息。此外,得益于Mamba的高效序列建模,这个解码器在构建实时语义分割网络方面也显示出巨大的潜力。
  3. 作者在三个广泛使用的遥感图像语义分割数据集上进行了综合实验。结果显示,作者的PyramidMamba与基于CNN和Transformers的最先进方法相比,取得了具有竞争力的准确性。

II Related Work

CNN-based Semantic Segmentation

远程感知图像理解的基本解释工具是语义分割。在过去的十年中,卷积神经网络(CNN)利用其分层结构、自动特征学习和端到端的方式,在远程感知图像的语义分割领域占据主导地位[21, 22, 23, 24]。全卷积网络(FCN)[12]是第一个基于CNN的端到端分割网络,标志着语义分割领域的重大进展。然而,FCN的过于简化的全连接解码器常常导致分割图粗糙。

为了解决这一挑战,开发了对称的编码器-解码器架构[25]。编码器逐步减少图像的空间维度,同时增加通道数以捕获高级语义特征,而解码器逐渐恢复空间维度并强化细节表示。最著名的网络是U-Net[13]及其变体[26],它们有效地缓解了粗糙分割问题,并保持了地理目标的丰富细节。尽管U-Net系列的结果有了很大的改进,但在处理复杂的远程感知场景时仍然面临挑战。这些网络的有限的局部感受野限制了它们捕获全局上下文信息的能力[27]。因此,这些网络难以挖掘地理目标之间的关键空间依赖关系,并提高对远程感知图像准确分割的全局理解能力。

Attention-based Global Context Modeling

为了解决传统卷积神经网络(CNNs)在遥感图像语义分割中的局限性,一些研究引入了注意力机制作为加强CNN全局上下文建模的关键技术。DANet [15] 提出了一种双重注意力机制,包括通道注意力与空间注意力,以同时捕获这两个维度的全局依赖关系。CCNet [28] 开发了一种交叉注意力块,可以通过交叉特征融合捕获密集的全局上下文信息。其他一些研究试图通过增加卷积核大小[29]或合并多尺度语义特征[30]来扩大感受野。特别是,多尺度特征融合已被证明是提高CNN性能并获取细粒度分割结果的有效途径。著名的PSPNet [18] 提出了一种空间金字塔池化模块,以提取并合并多尺度语义特征,并在语义分割领域取得了重大突破。然而,通过空间池化和上采样操作提取的多尺度特征存在同质语义信息,削弱了特征融合的有效性。此外,上述两种方案仍然过于依赖卷积操作,并未真正摆脱局部模式的限制。

在近两年中,视觉 Transformer (ViTs)[17]将2D图像解释视为1D序列建模,逐渐成为计算机视觉任务尤其是语义分割[31]的主流方法。与基于注意力的CNN相比,ViTs采用纯粹的自注意力结构,展现出更强大的全局上下文建模能力。尽管ViTs在全球上下文信息提取方面表现出色,但它们在计算效率和局部特征表示方面存在不足。为了提高ViTs的效率,一些研究专注于设计层次结构[32]或开发有效的注意力机制,如基于窗口的注意力[33]和线性注意力[34]。至于提高局部特征表示,最常见的方法是将通过CNN提取的局部特征与通过ViT提取的全局特征进行融合[35]。这种方案虽然有效提高了语义信息,但未能实现更细粒度的多尺度特征表示。

Vision Mamba

基于上述分析,基于CNN和基于ViT的方法在扩大感受野方面都有其缺点。使用多个感受野的CNN会导致多尺度特征融合中的同质信息冗余。而应用ViT进行全局上下文建模则表现出较低的效率。最近,一种基于选择状态空间模型(SSM)[20]的新型架构——Mamba,为视觉理解开辟了新途径。Mamba独特地采用变量参数来表示全局依赖,并利用硬件优化的计算策略来平衡内存效率和性能。此外,Mamba的选择性扫描机制使其能够专注于挖掘长序列的核心语义,从而解决语义冗余问题。得益于这些独特特性,许多研究者已成功将Mamba应用于计算机视觉任务[36],如图像分类[37]和语义分割[38]。在遥感领域,Mamba已被探索用于密集预测任务,例如变化检测[39]和语义分割[40],并与CNN和ViT相比获得了显著改进。

受到Mamba的选择性扫描机制和高效率的启发,作者将它作为一种金字塔特征融合的连接器,从而进一步增强多尺度特征表示。具体来说,作者将Mamba与一个空间金字塔模块集成,该模块应用更多的池化尺度以产生更丰富的多尺度语义内容。这种创新组合实现了无冗余特征融合和高效率的双赢局面。

III Methodology

Preliminaries

在本节中,作者简要介绍了状态空间模型(SSM)的原理,这是Mamba的基础。更详细的内容可以在原论文[20]中找到。

图1:作者的PyramidMamba概览。(a)网络结构,(b)密集空间金字塔池化,(c)Mamba块。

Mamba块继承了上述优点,应用于作者的PyramidMamba中。如图1和2所示,选择扫描机制首先将金字塔序列扩展为四个双生序列。然后,将选择路由和S6块[20]应用于压缩四个双生序列并从每个序列中提取核心语义信息。最后,使用合并操作产生输出。

Overview Architecture

在本节中,作者主要介绍了PyramidMamba的结构,它是基于经典的编码器-解码器架构,如图1所示。输入图像被送入图像编码器以提取高级特征和低级细节特征。然后,基于Mamba的解码器处理高级特征以增强多尺度语义内容。最后,处理过的高级特征与低级细节特征融合,以增强分割结果的空间细节。在接下来的章节中,将详细描述基于Mamba的解码器的组成部分,即密集空间金字塔池化和金字塔融合Mamba。

Image Encoder

对于细粒度语义分割,同时保留低级细节和高级语义信息是非常关键的。在作者的PyramidMamba中,作者引入了一个分层图像编码器,从输入的遥感图像中提取低级详细特征和高级语义特征。具体来说,作者使用轻量级CNN(ResNet18)[41]和基于窗口的ViT(Swin-Base)[33]作为编码器。因此,PyramidMamba有两个版本。与ResNet18的结合构建了一个轻量级的分割网络,适用于实时应用。同时,采用Swin-base构建了一个大型分割模型,以产生更精确的分割结果。此外,图像编码器的轻松切换可以说明基于Mamba的解码器的实用性。

Mamba-based Decoder

为了解决多尺度特征融合中的信息冗余问题,作者设计了一个基于Mamba的解码器。具体来说,作者首先开发了一种密集的空间金字塔池化方法来获取富含多尺度语义内容的特点图。然后,作者采用Mamba块,它可以使用其选择性的过滤机制来减少多尺度特征图中的语义冗余。最后,融入了一个卷积前馈神经网络以进一步增强多尺度特征表示。

Ii-D1 Dense spatial pyramid pooling (DSPP)

Ii-D2 Pyramid fusion Mamba (PFM)

因此,作者部署了标准的Mamba块进一步处理多尺度特征,如图1(a)和(c)所示,它可以使用自身的选择性过滤机制有效地表征跨尺度的核心语义。作者首先应用展平操作从多尺度特征生成金字塔序列:

通过遵循这种设计,基于Mamba的解码器有效聚合多尺度特征,减少信息冗余,并增强多尺度特征表示,以实现精细的语义分割。

Loss Function

IV Experimental Settings and Datasets

以下是IV 实验设置与数据集部分的开始:

在本节中,作者将介绍实验设置以及所使用的数据集。

Datasets

为了评估所提出的PyramidMamba的性能,使用了三个公开可用的遥感语义分割数据集进行实验,包括OpenEarthMap数据集[42],ISPRS Vaihingen数据集和ISPRS Potsdam数据集。以下是这些数据集的详细信息。

Iv-A1 OpenEarthMap

OpenEarthMap数据集是一个大规模的高分辨率土地覆盖映射数据集,它由5000张图像组成,包含八种土地覆盖类别(裸地、草原、开发空间、道路、树木、水体、农业用地、建筑物)。图像的空间分辨率在0.25米到0.5米之间。空间分布覆盖了来自六大洲44个国家的97个区域。由于OpenEarthMap数据集具有广泛的空间变异性、复杂的地理目标和场景,对其进行语义分割非常具有挑战性。在OpenearthMap数据集中,每个区域的遥感图像被随机分为训练集、验证集和测试集,分别产生了3000、500和1500张图像。在作者的实验中,由于测试集不公开,作者使用验证集进行定量比较。输入图像被统一调整为1024×1024像素的块,并且在训练和测试阶段使用了数据增强策略,如水平和垂直翻转。

Iv-A2 Vaihingen

Vaihingen数据集由33个细分辨率的图像块组成,平均大小为2494×2064像素。每个图像块包含三个多光谱波段(近红外、红色、绿色),以及一个数字表面模型(DSM)和归一化数字表面模型(NDSM),其地面采样距离(GSD)为9厘米。该数据集包含五个前景类别(不透水表面、建筑物、低矮植被、树木、汽车)和一个背景类别(杂乱)。在实验中,仅使用了图像块。图像块被裁剪成1024×1024像素的块。在训练模型时,使用了包括水平垂直翻转、随机缩放和裁剪,以及随机马赛克在内的数据增强策略。

Iv-A3 Potsdam

波茨坦数据集包含38张超高分辨率航拍图像(地面采样距离5厘米),尺寸为6000×6000像素,涉及6个地理目标类别(不透水表面、低矮植被、树木、汽车、建筑物和杂物),四个光谱波段(红、绿、蓝和近红外),以及数字表面模型(DSM)和归一化数字表面模型(NDSM)。在实验中,作者遵循官方的训练和测试划分,并且只使用了三个波段(红、绿、蓝)。原始图像块被裁剪成1024×1024像素的 Patch 作为输入,作者采用了随机翻转和随机马赛克作为数据增强。

Evaluation Metrics

作者使用整体准确度(OA)、平均交并比(mIoU)、F1分数、精确度和召回率来评估模型的性能,定义如下:

Experimental Setting

所有实验中的深度模型都是使用PyTorch框架在单个NVIDIA GTX 4090 GPU上实现的。采用了AdamW优化器来训练深度模型。使用了多项式学习率调整策略,指数参数设置为0.9。基础学习率设置为6e-4,而图像编码器的学习率特别设置为6e-5。批量大小和权重衰减分别设置为2和0.01。总训练周期设置为45,并且在最初的5个周期内应用了 Warm up 训练策略。使用了早停策略来防止过拟合。在测试阶段,作者应用了数据增强技术,如水平翻转和垂直翻转以及多尺度变换,这也被称为测试时增强(TTA)。

Benchmark Methods

为了验证所提出方法的有效性,作者选择了一系列最先进的分割方法进行全面的比较,包括:1)实时语义分割网络:BiSeNet [43],ShellNet [44],SwiftNet [45],ABCNet [46],和UNetFormer [47];2)基于CNN的语义分割网络:U-Net [13],PSPNet [18],DeepLabV3+ [30],DANet [15],UFMG-4 [48],ResUNet-a [22],MANet [27],LANet [21],DDCM-Net [49],EuNet [50];3)基于Transformer的语义分割网络:SegFormer [51],Segmenter [31],SwinUperNet [33],BoTNet [52],DC-Swin [53],SwinB-CNN [54],CG-Swin [55],Mask2Former [56];4)基于Mamba的语义分割网络:RSMamba [57];5)用于语义分割的视觉语言模型:CLIPSeg [58];6)最近的遥感图像分割网络:FTransUNet [59],SAPNet [60]和MMT [61]。

V Experimental Results and Analysis

Ablation Study

为了验证所提出模块的有效性,作者在ISPRS Vaihingen数据集上进行了消融实验。为了确保消融实验的公平性,在测试阶段作者没有应用任何测试时间增强,并且图像编码器统一设置为ResNet18。

V-A1 Network variants

如表1所示,Baseline包括图像编码器和上采样操作。BaselineDSPP表示Baseline与密集空间金字塔池化的结合,而BaselineDSPPFM代表不包括低层次细节特征的整个网络。

V-A2 The effectiveness of each component

在提出的PyramidMamba中,DSPP通过简单的拼接操作为细粒度分割编码丰富的多尺度语义信息。如表格I所示,部署DSPP可以使mIoU提高3.5%,这可以说明其在多尺度特征表示中的有效性。此外,PFM的利用可以进一步将mIoU提升1.1%。这一结果不仅展示了PFM的有效性,也表明了Mamba在多尺度特征融合中的显著优势。

V-A3 The superiority of the dense spatial pyramid pooling (DSPP)

如表2所示,作者将所提出的DSPP与PSPNet中的标准空间金字塔池化(SPP)模块进行了比较。选择Baseline作为基本网络。结果显示,与SPP相比,作者的DSPP在mIoU上提高了1.4%,在F1得分上提高了0.9%,这可以说明密集池化在加强多尺度表示方面的优势。

V-A4 The effectiveness of aggregating the low-level detailed feature (LDF)

引入空间详细特征是优化语义分割结果的有效方法。由于具有更高分辨率,分层深度网络中的低级特征包含了丰富的空间细节。为了展示低级详细特征对准确度的贡献,作者在消融实验中将其移除。如表3所示,采用低级详细特征可以将mIoU指标和F1分数分别提高2.7%和2.1%,证明了其有效性和必要性。

实时语义分割。值得注意的是,在本节中为了公平比较,使用了测试时间增强。如表4所示,网络的速度(FPS)是通过单个NVIDIA GTX 4090 GPU上的两个1024×1024图像块来测量的。结果显示,与其它先进的实时分割网络相比,所提出的PyramidMamba在保持竞争力速度的同时,在准确度上具有优势。特别是,与最近的实时ViT(UNetFormer)相比,作者的PyramidMamba在 indoU 上提高了0.4%,并且至少比其它实时CNNs在mIoU上高出1.8%。这些结果不仅证明了作者PyramidMamba的优越性,还说明了基于Mamba的方法在构建实时深度网络中的巨大潜力。

Quantitative Comparisons with State-of-the-art Semantic Segmentation Methods

为了进一步验证所提出方法的有效性,作者将其与最先进的方法在三个公开可用的数据集上进行比较,即OpenEarthMap数据集、ISPRS Vaihingen和Potsdam数据集。此外,为了确保公平比较并展示作者基于Mambo的解码器的适用性,图像编码器被设置为广泛使用的ViT,即Swin-Base。

V-C1 OpenEarthMap

OpenEarthMap数据集包含许多复杂场景和令人困惑的地物目标。因此,在这个数据集上实现高精度是非常具有挑战性的。如表5所示,作者的PyramidMamba实现了70.8%的mIoU,分别比基于CNN的方法MANet和基于ViT的方法SegFormer高出6.8%和4.8%的mIoU。值得注意的是,作者的方法在特定类别上也取得了最高的准确度,例如IoU-Road(64.9%)、IoU-Building(79.6%)和IoU-Developed(57.9%)。这些多尺度地物目标上的显著改进可以证明作者的PyramidMamba在多尺度特征表示上的有效性和优越性。可视化结果可以进一步证实这一点。如图2所示,与UNet和UNetFormer相比,提出的方法不仅能够分割具有精细形状的建筑(第一行),还能保持道路的连续性(第二行)。此外,对于容易混淆的开发土地,作者的PyramidMamba也有显著优势。

Iv-A2 Vaihingen 2024-06-30-04-18-49

国际摄影测量与遥感学会(ISPRS)的Vaihingen数据集是验证遥感图像语义分割方法有效性的广泛使用的数据集。已有数百个深度模型被开发出来,并在该数据集上取得了高分。因此,要获得进一步的准确性突破是困难的。然而,正如表6所示,作者的PyramidMamba获得了最佳的mIoU(84.8%)和整体准确性(93.7%)。特别是,作者的PyramidMamba在OA上比最近的遥感图像分割方法SAPNet和MMT高出4.0%,在mIoU上高出0.7%,这可以证明作者方法的前进和优越性。此外,对于采用相同Mamba基础架构的RSMamba,作者的PyramidMamba在mIoU上增加了1.9%。同时,可视化结果进一步见证了作者的方法的优势。如图3所示,得益于细粒度的多尺度特征融合,作者的PyramidMamba可以确保建筑物分割的完整性,尽管建筑物表面非常复杂。对于小型的地理目标汽车,作者的方法也能保持良好的分割形状。#### V-B3 波茨坦

图3:在Vaihingen数据集上的可视化比较。

国际摄影测量与遥感学会的Potsdam数据集也是用于遥感图像语义分割的广泛使用的数据集。在这个数据集上,作者的PyramidMamba获得了最先进的mIoU(88.0%)、平均F1分数(93.5%)和IoU-Car(96.9%),在mIoU上至少比基于CNN的方法高出1.1%,在平均F1分数上比基于Transformer的方法高出0.3%以上。这些结果进一步证明了作者PyramidMamba的有效性和优越性。作者还与基于Transformer的方法CG-Swin进行了可视化比较。如图4所示,作者的PyramidMamba在检测狭窄道路方面具有显著优势。

VI 结论

在本文中,作者提出了一种新型的基于Mamba的解码器,用于遥感图像的语义分割,名为PyramidMamba。为了解决多尺度特征融合中的语义冗余问题,作者在解码器中引入了标准的Mamba块,利用其选择扫描机制来增强多尺度特征表示。

此外,作者提出了一种密集空间金字塔池化来实现细粒度的金字塔特征。得益于以上方法,作者的PyramidMamba在三个公开可获得且广泛使用的遥感图像分割数据集上与最先进的方法相比展示了优势。

同时,消融研究也说明了所提出解码器中每个组件的有效性,并揭示了其在构建实时语义分割网络方面的巨大潜力。在未来,作者将继续探索Mamba基础结构在多模态学习和基础模型中的潜力。

参考

[1].PyramidMamba: Rethinking Pyramid Feature Fusion with Selective Space State Model for Semantic Segmentation of Remote Sensing Imagery.

# MambaIR

性能超越SwinIR!MambaIR: 基于Mamba的图像复原基准模型 , 首次将状态空间模型引入到图像复原中。

这里介绍一下我们近期使用状态空间模型Mamba来做图像复原的工作。

论文和代码都已开源,欢迎感兴趣的小伙伴关注

[论文链接]

https//arxiv.org/pdf/2402.15648.pdf

[代码链接]

https//github.com/csguoh/MambaIR

我们这个工作首次将状态空间模型引入到图像复原中。由于之前很少有相关的探索,因此为了更加深入理解,这里需要回答以下两个问题:

  1. 为什么Mamba可以用来做图像复原?

首先我们来看一下状态空间模型的迭代表达式(这里直接把最终离散化后的结果给出了,详细推导可以参见Mamba的论文)

可以看到,状态空间模型采用了一种迭代的形式,首先从第0个时间步开始,按照上述公式进行迭代来得到之后时间步的预测。读到这里,大家可能会发现状态空间模型的表达式好像和RNN有点像啊。而从近几年的图像复原的模型中可以看到,很少有使用RNN来做图像复原主干的,因此我们很自然想到,使用Mamba(也是状态空间模型的一种)来做图像复原make sense吗?

实际上,上述迭代式是有特殊结构的,并在数学上可以化为卷积的形式,也就是说状态空间模型可以同时理解为RNN和CNN,而CNN对于图像复原则是很常见的模块,这也就为Mamba在图像复原中的应用奠定了基础。下面我们看一下将状态空间模型推导为CNN卷积形式的证明,更详细的证明可以看下Mamba的原始论文。

以上就是状态空间模型可以转化为卷积的证明,这也在理论上为Mamba在图像复原中的未来探索提供了支持。

当然在实际上我们的工作MambaIR的工作也是在此基础上进行了改进,包括

  • 使用Mamba中的策略——由输入生成卷积权重(类似于动态卷积,所以表达能力还是很强的)
  • 使用四个方向的扫描,让当前像素的输出实际上是来自上下左右四个方向邻域的一维卷积结果
  1. 既然可以做,那么效果怎么样呢?

首先谈一下个人感受,我们提出的MambaIR是一个很简单的Baseline,并没有引入十分复杂的设计,对比同样简单的基于Transformer的图像复原基准模型SwinIR,我们的MambaIR是要比SwinIR明显要高的(见下图)。因此,不难想见,未来针对Mamba做特定设计将会进一步提升基于Mambda的图像复原模型的能力。

经典图像超分x2倍率实验结果

动机

Transformer模型已经在底层视觉领域取得了十分瞩目的成果,例如IPT,SwinIR等等。最近随着先进的状态空间模型Mamba的出现以及其在NLP中的喜人表现,使用Mamba来做各类视觉任务变得十分有趣。

对于图像复原任务来说,我们总结出Mamba相较于之前的CNN和Transformer具有如下优势:

  • 对于CNN来说,CNN具有静态权重以及局部感受野,这大大限制了模型利用更多像素的能力,正如HAT指出的那样。而Mamba则具有全局的感受野,可以以整个图像作为token序列作为输入。
  • 对于Transformer来说,虽然标准的注意力机制具有全局感受野,但是在图像复原任务中使用这一机制将会带来无法接受的计算代价,因此为了折中,目前的工作大多使用了SwinTR的机制,但是这同样限制了感受野同时window的边缘也容易产生伪影。

下面是MambaIR与基于CNN的方法(EDSR,RCAN)和基于Transformer的方法(SwinIR,HAT)在有效感受野上的对比结果,可以看到MambaIR具有正幅图像范围的感受野,从而可以更加充分地利用图像块重复先验。

有效感受野对比

方法

首先需要指出的是,我们MambaIR的主要目的是提出一个简单的baseline模型,并方便之后的进一步研究改进。因此在模型设计上,我们遵循了之前图像复原任务的大致流程,即浅层特征提取,深层特征提取,以及高质量图像重建。

这里简单介绍MambaIR的核心组建,残差状态空间模块(Residual State-Space Block,RSSB)

对于特征X,其首先经过LayerNorm以及视觉状态空间模块(VSSM)来建模空间维度上的远程依赖,并伴随着可学习的scale参数来调整来自残差连接中的权重:

其中VSSM模块通过使用四个方向的扫描来将2D图像转化为1D输入序列,并使用前面的离散化的迭代公式进行建模。关于该模块的具体细节可以参考原论文。

此外,在实验中我们同样也和原始的Mamba的作者那样,发现在Mamba后面加MLP甚至会带来性能的下降。为了进一步补充标准的Mamba的建模能力,我们进一步考虑了两个图像复原的特定先验:

  • 局部邻域重复:我们使用卷积来补充邻域像素的相似性
  • 通道交互:我们引入通道注意力来补充空间维度的交互

最终实验结果也发现我们引入的这两种先验比标准的Mamba性能更好,比简单地使用MLP的性能也更好。具体消融实验结果如下:

结语

尽管目前针对Mamba的各类任务模型已经在不断涌现,但是目前这一任务仍然存在着较大的改进空间。相信未来随着大家对Mamba认知的加深,可以有更多更好的工作出现。我们MambaIR的工作还在持续更新中,之后会包含更多的复原任务,欢迎大家持续关注。如果有任何问题或者进一步的想法,也欢迎随时讨论。


​​​​​​​m合集~ ~ ~~~ 天皓智联 whaosoft aiot http://143ai.com

标签: 人工智能

本文转载自: https://blog.csdn.net/qq_29788741/article/details/140191520
版权归原作者 whaosoft143 所有, 如有侵权,请联系我们删除。

“Mamba~合集1”的评论:

还没有评论