FusionMamba: Efficient Image Fusion with State Space Model【文献阅读】

论文：FusionMamba：一种基于SSM的有效图像融合方法
arXiv：https://arxiv.org/abs/2404.07932
作者单位：中国科学院自动化研究所、模式识别重点实验室、电子科技大学
推荐阅读：深入浅出一文图解Vision Mamba

Abstract

图像融合的目的是生成高分辨率的多光谱/高光谱图像（multi/hyper-spectral image），通过结合有限空间信息的高分辨率图像和具有大量空间信息的低分辨率图像。当前基于深度学习的图像融合方法主要是基于CNN和Transformer模型去提取图像的特征和融合不同类型的数据。CNN虽然和有效，但是限制于它的感受野，限制了他们去进行全局建模的能力。Transformer虽然能够专注于全局建模，但被它的二次计算复杂度所阻碍。然而，最近的SSM模型相关的研究，尤其是Mamba，不仅仅能够进行全局建模还能够保持线性复杂度。在信息融合中有很少的研究在探索SSM的潜力，探索在图像融合领域上的关键能力。因此。我们提出的FusionMamba模型，一个创新性的方法用于图像融合。我们的贡献主要集中在2个方面。第一点：不同的源图像有着不同的属性和差异，我买的Mamba模块包含2个共享U型网络，呈现出一种创新性的结构用于高效、独立、分层的来提取图像的空间和光谱特征。第二点：为了有效的结合空间和光谱信息，我们扩展了Mamba模块去容纳2个双分支的输入。这样的扩充创建了一个新的模块叫做FusionMamba快，优于现存的拼接和交叉注意力的方法。为了验证FusionMamba模块的有效性，我们构建了一系列数据集在5个相关的数据集在三个融合任务上。从定性和定量的角度证明了我们的方法实现了SOTA的性能，强调了FusionMamba的优越性。

1 Introduction

图1：结合不同网络组建在性能和效率上的对比。特征提取模块包含卷积，自注意力，Mamba模块。另外，信息融合的方法包含Concat和交叉注意力，以及提出的FMamba块。在相同的参数计算下，我们的方法在合理的FLOP下拥有较好的性能。使用的数据集是WV3。

如图1所示，实验结果证明了FusionMamba模块比现存的融合技术更好，比如拼接和交叉注意力。总的来说本文贡献如下：

• 我们的Mamba模块包含2个U型的网络，称作空间U-Net和光谱U-Net。这个结构设计的目的在于实现高效、独立、层次的学习空间和光谱特征，让它适用于双分支的图像融合任务。
• 为了整合空间和光谱信息，我们扩展这个Mamba模块去处理双分支输入，创建了一个创新性的模块叫做FusionMamba模块。对于现存的技术来说，这个模块表明了较好的有效性，这是一个成功的扩展将SSM到信息融合领域。
• 提出的方法证明了SOTA的性能，在5个数据集上的三个不同图像融合任务。这些结果表明了FusionMamba的有效性。

2 Related Works

2.1 DL-Based Methods for Image Fusion

省略…较常见，下面基于SSM的相关工作可以一起学习一下。

2.2 State Space Model

SSM模型是一种基础科学模型，主要用于控制理论。最近在应用在深度学习领域，最近的研究有LSSL和S4。LSSL接近SSM作为一个基础的深度学习模型通过一系列数学相关的发展。另外一个方面，S4引入了正常+低秩的概念，从而有效降低了SSM的计算复杂度。随后的工作，比如S5和H3，进一步的调查了SSM在深度学习中的潜力。缩小了SSM和Transformer的差距。这些努力导致了Mamba的创建，它融合了先前研究的关键以及提出了一个选择性的机制用于动态提取特征从输入的数据。Mamba在各种1D的数据集上优于Transformer并需要更少的计算资源。Mamba的成功吸引了计算机视觉研究者的注意，使其在各种视觉任务上有一个广泛的应用。值得注意的贡献包括：Vision Mamba和VMamba以及Pan-Mamba。Mamba最开始的设计是用于1D任务以一个简单的单方向，Vision Mamba引入了多方向的序列建模去增强空间感知上的理解。VMamba通过交叉扫描进一步的提升了这一点，一种4个方向的建模方法，能够发觉额外的空间联系。Pan-Mamba代表一个开创新的研究利用Mamba去解决这个全彩图像增强，基于Mamba模块产生了较好的结果。上述方法主要集中在Mamba的应用性和方向性。然而，这里还有部分尝试在于扩展SSM的潜力去整合不同类型的信息，这在图像融合领域是很重要的。

2.3 Motivation

现存基于深度学习的图像融合方法主要应用CNN和Transformer来进行特征提取和信息融合。尽管CNN和高效，但是很难去捕获全局信息。相反的，Transformer利用全局的感知能力，但它的计算复杂度较高。幸运的是，最近提出的Mamba模块好像能够解决这样的窘境。然而，直接堆叠 Mamba模块去进行特征提取并不可取。受到U2Net的鼓舞，我们整合了Mamba模块到2个不同的U型网络：空间U型网络和光谱空间网络。这样的策略能够有效的学习特定和分层的空间和光谱特征。除开特征提取，信息整合也是图像融合的关键。不幸的是，只有很少的研究在探究SSM在整合不同类型信息的潜力。因此，我们扩展了单输入的Mamba块去支持两个输入，提出了一个创新性的模块叫做FusionMamba块，能够有效的整合空间和光谱信息。

3 Method

3.1 Notations

PAN/RGB图像被定义为P（H，W，S）H和W及S分别代表图像的高、宽、通道数。M（h，w，S）代表LRMS/LRHS图像，S代表光谱带。其中h和w是H的四分之一。此外，上采样的LRMS和LRHS图像，生成HRMS和HRHS图像和这个真实标签被分别定义为M_U、O_~及O分辨率都是（H，W，S）。其中Fia和Fib代表第i个阶段的空间和光谱特征图，一共有5个阶段。其中F1和5的维度都是（H，W，C），F2和4是二分之一的H和W，维度是2C。其中F3是四分之一的H和W，维度是4C，其中C代表通道维度。另外，N代表SSM隐藏状态的大小。

3.2 Preliminaries（准备工作）

State Space Model.：SSM是一个连续的方法能够将1D的输入x映射成y，通过中间的隐藏状态h。数学上，它通常能够用下面的常微分方程表示，如：

A是一个演化参数，B和C是一个投影参数。这个等式表现出SSM展现出的全局感知能力，当前的输出受到所有先前输入的影响。当A、B、C是常量时，公式1定义为一个线性时间不变的系统LTI Linear time-invariant，这就是在S4中的情况。另外，一个描述linear time-varying LTV系统，作为Mamba【参考文献11】。LTI固有缺乏感知输入内容的能力，但是LTV的设计具有这样的能力。关键的差异就是Mamba补充了SSM的限制。

Discretization.：（离散化）当SSM被应用在深度学习上，它需要被离散化。为了实现这个目的，一个时间刻度的参数被定义为△，把连续的A和B变为离散的A-和B-。通过引用ZOH zero-oder hold零阶保持器作为转换的算法，离散的参数被下式定义。

Selective Scan.：在Mamba模型中，随着输入的不同导致参数不同，公式3不能改写为卷积的形式，这阻碍了SSM的并行化。为了应对这样的机制，Mamba引入了选择扫描机制，整合了三种经典的技术：核融合、并行扫描和重新计算。通过选择扫描的算法，Mamba实现了较好的运行速度并且仅需要较低的内存开销。

3.3 Network Architecture

为了完全利用SSM的潜力，受到参考文献29U2Net的鼓舞，我们整合了Mamba 模块作为其网络结构。这个结构能够独立和分层的学习不同类型的信息，主要是通过2个U-型的网络：空间U-Net和光谱U-Net。前者从P中提取空间细节，后者从M中捕获光谱特征。为了能够充分的捕获深层信息，并不显著的增加网络参数，我们提取三个不同尺度的特征。这代表每个U-Net包含一共5个阶段。每个阶段，特征图最初都要经历一个Mamba模块。随后，将其传入一个FusionMamba模块用于信息的整合。不同类型的卷积层被用于上下采样，以及通道维度的调整。M在生成O~上扮演更加重要的角色相比于P，FusionMamba块能够将其特定的融入光谱U-Net去模拟逐步注入空间信息到光谱特征图中。

3.4 Mamba and FusionMamba Blocks

我们网络的Mamba和FusionMamba块的范式如图3所示。为了增强空间感知能力，我们采用参考文献26中的4个方向序列建模方法，如图的右侧所示。这部分提供了Mamba模块和提出的FusionMamba模块的细节，伴随着FLOPs的分析。

图3：左边是Mamba模块的原理图，中间是提出的FusionMamba模块，右边是4个方向图解。这个Mamba模块是单输入的特征提取模块，FusionMamba模块是一个双输入的结构被设计来整合不同类型的信息。

Mamba Block.：作为一个输入的特征Fin H，W，C，我们最开始用LN层来对其进行归一化。随后，用并行的1×1卷积层来处理，产生2个不同的特征，分别定义为X和Z H，W，C，这个处理可以用下式表达。

之后，X被4个方向展平处理为x1、x2、x3和x4，每个的维度是HW×C。这些1维的序列通过SSM模块处理用于特征提取，结果得到4个输出：

算法1提供了SSM模块的细节，Linear代表全连接层。接下来，我们把SSM的输出unflatten回去，把它们结合起来去得到一个新的特征图定义为Y H，W，C。然后特征图通过Z门控gated得到。结果tensor经过1×1卷积层然后加到输入Fin上，最终的输出用Fout H，W，C表示。

SiLU代表激活函数，简化描述，定义上式为：

这里的Mamba定义为Mamba模块。

FusionMamba Block.：原始的SSM仅仅支持单个输入，为了有效的整合不同类型的信息，我们扩展它去使用双分支输入，产生了FSSM模块，它的细节如算法2。在FSSM模块中，一个输入代表去生成映射和时间刻度参数。另外一个输入是要处理的序列。包含8个FSSM块，融合FusionMamba模块被设计成一个同步对称的结构。输入的空间和光谱特征分别为Fa_in和Fb_in，H，W，C，我们在Mamba中去生成两组特征图如下所示：

这个等式是公式4的直接推广，符合的解释就先省略了。下一步把Xa和Xb直接flatten成4个方向。1D序列随后传入FSSM块进行信息的整合。

FSSMa和FSSMb代表图3中FusionMamba模块的左边和右边分支。之后，我们分别处理了这两组输出，产生了2个新的特征图，定义为Ya和Yb维度为H，W，C。这些特征图随和结合起来的到输出Fout：

将上式简化如下所示：

其中的FMamba为FusionMamba模块。

基于SSM的算法1和算法2如下所示：

Analysis of Flops.：具有D参数的卷积层，FLOP的计算被记作2HWD。Mamba模块的参数为D，选择扫描花费的参数为9HWCN FLops，总的FLOP为 2HWD+36HWCN。进一步来讲，FusionMamba模块有相同的参数消耗：2HWD+72HWCN FLOPs。自注意力和交叉注意力，总的FLOP的估计为2HWD+4HHWWC FLOPs。这些方法的对比描述如图4。对比了卷积、Mamba模块、FusionMamba和自/交叉注意力的FLOP。注意，虽然卷积有较低的FLOPs，但是它缺乏捕获全局信息的能力。

3.5 Loss Function

本文的主要贡献在于SSM的应用与创新，仅仅使用L1来训练网络即可：

P、M和O分别代表训练集中的PAN（Pansharpening 全色图像，灰度图，全波段，高分辨率）/RGB、LRMS（Low-Resolution Multispectral低分辨率多光谱）/LRHS（Low-Resolution Hyperspectral低分辨率高光谱）图像，以及GT图像。f代表我们网络，其中seta代表可学习的参数，T是总的样本数量。另外，使用的是L1范数作为损失。

4 Experiment

4.1 Pansharpening

Dataset.：对于全色图像增强任务，我们主要在WorldView-3 WV3数据集上进行。这个数据集包含了10000张训练样本，百分之90作为训练集，剩下的作为验证集。另外，它包含了20个低分辨率和全分辨率的测试样本。每个训练样本包含一个三元组的格式：PAN/LRMS/GT，分辨率分别是：64×64×1、16×16×8和64×64×8。测试样本也包括PAN/LRMS/GT三元组图像集，分辨率分别为：256×256×1，64×64×8，256×256×8。另外，全分辨率的图像对的PAN和/LRMS图像对分辨率为：512×512×1和128×128×8。本节使用的数据集来自参考文献4、更多的细节可以查看补充材料。

Benchmarks and Evaluation Metrics.：我们比较了最近的SOTA工作，包括两种传统算法[36] [33]和8个基于深度学习的算法[40] [46] [3] [38] [23] [24] [29] [19]。为了确保公平性。基于深度学习的方法用相同的3090显卡和PyTorch环境。按照全色图像增强的标注，我们应用4个质量PSNR、Q8、SAM和ERGAS指标。去评估低-分辨率实验结果。对于全分辨率，选择Dλ、Ds和QNR作为定量指标。

Network and Training Settings.：在全色增强任务我们设置C和N为32和8。对于训练，epoch、batch size和初始化学习率分别设为400、32和5×10-4。另外，我们利用Adam优化器并每200轮把学习率减半。基于其他深度学习的方法我们用其默认参数做处理，更多信息可以参考补充材料。

Evaluation Results.：实验结果如下所示。

4.2 Hyper-Spectral Pansharpening

Dataset.：对于高光谱全色增强任务，我们在3个流行的数据集上进行：Pavia、Botswana和WDC。
内容和前面的数据描述差不多，省略…

Benchmarks and Evaluation Metrics.：省略…

Network and Training Settings.：对于高光谱全色锐化，我们设置C为48，N为8。进一步讲，U型网络，特征的通道数保持不变，与图3的描述有点儿偏差。训练上述3个数据集时，epoch、batch size 和初始化学习率分别设为2000/2500/2000、32、2×10-4，还是选择Adam优化器学习率每1000轮减半。

Evaluation Results.：定量和定性的结果如表2和图6，明显的FusionMamba优于其他模型的性能。

4.3 HISR

相关描述省略，实验结果如下所。

4.4 Ablation Study

Effectiveness of Network Design.：为了验证我们网络设计的有效性，我们设计了6种FusionMamba的变体。V1不使用U型网络用语特征提取，V2结合空间和频谱U-Net网络为单个U-Net。V1和V2被设计来展示U-型结构的有效性。V3和V4取代Mamba模块和FusionMamba模块种4个方向的序列建模为单向和双向序列建模。这些变换能够验证不同角度flatten方法的效果。V5直接将Fa_out作为FusionMamba模块的输出，V6直接将Fb_out作为输出。这些变体是为了评估FusionMamba模块的有效性来设计的。我们在低-分辨率的WV3数据集上来验证这些方法的有效性。为了确保公平性，我们确保所有网络的参数是相同的。结果如表4清晰的证明了我们网络设计的有效性。

Effectiveness of Mamba and FusionMamba Block.：为了验证Mamba模块特征提取的有效性，我们用卷积Conv，自注意力SA替代它。为了验证信息融合FusionMamba在信息融合上的有效性，我们用Concat和cross-attention来替代它。这些共组成6种不同的变体。在WV3低分辨率图像的数据集上来验证它的性能，并保证他们拥有相同的参数。定量实验如图5所示，清晰的证明了我们提出的方法实现了最好的性能以及拥有一个较为理想的FLOPs。

5 Conclusion

本文FusionMamba是一个具有创新的高效图像融合方法。为了有效的提取不同特征，我们整合了Mamba模块到U型的网络中，包含：空间U-Net和光谱U-Net。这个方法能够有效充分并分层的学习空间和光谱特征特征。进一步的为了实现复杂的信息融合，我们扩展Mamba模块为双分支输入，提出了一个新的模块称作FusionMamba模块。我们的FUsionMamba实现了SOTA的性能在5个数据集的3个融合任务，突出了我们方法的有效性。

总结

2024年05月02日11:30:02

1. 将Mamba模块基于U-Net的Baseline网络嵌入到了多/高光谱全色图像增强任务中，还提出了能够进行双分支输入的FusingMamba模块；
2. 这个双输入，就是将两个输入放到Mamba模块中然后加起来？
3. Mamba模块基本理解了，其中的SSM还得再下来看看。
4. 整体感觉讲的还是很清楚的，但是并不是特别惊艳。Mamba系列值得自己持续学习，是否能够平替Transformer或者CNN，除开在参数量和计算复杂度的优势下，是否还能够达到更好的效果呢？继续学习吧！