随机 Shuffle 与逆 Shuffle | Shuffle Mamba 基于贝叶斯思想的随机扫描策略在多模态图像融合 !

图像处理大模型数据库

点击下方卡片,关注 「AI视界引擎」 公众号

( 添加时备注:方向+学校/公司+昵称/姓名 )

picture.image

picture.image

多模态图像融合集成来自不同模态的互补信息,产生增强和有用的图像。尽管状态空间模型,如Mamba,在用线性复杂度进行长程建模方面表现出色,但大多数基于Mamba的方法使用固定的扫描策略,这可能引入偏置先验信息。

为了减轻这个问题,作者提出了一种新颖的基于贝叶斯思想的扫描策略,即随机Shuffle,并由理论上可行的逆Shuffle来维持信息协调不变性,旨在消除固定序列扫描引起的偏见。

基于这个转换对,作者定制了Shuffle Mamba框架,深入考虑模态意识的信息表示和在空间和通道轴上的跨模态信息交互,以确保强大的交互性和多模态图像融合的公正全局感受野。

此外,作者还开发了一种基于蒙特卡洛平均的测试方法,以确保模型的输出与预期结果更接近。在多个多模态图像融合任务上的广泛实验表明了作者提出的方法的有效性,在同类最先进替代方案上获得了卓越的融合质量。

Introduction

多模态图像融合是计算机视觉中的基本任务,它涉及从不同成像模式捕获的同一场景图像中提取和集成有价值的信息。这一任务的目标是创建一个具有更全面和信息丰富的综合图像,典型应用包括超分辨率以及医学图像融合(MIF)。在超分辨率中,卫星受传感器限制,只能 capturing low-resolution multi-spectral (LRMS) 和pan-chromatic (PAN)图像。具体来说,PAN图像具有优越的空间细节,但光谱分辨率有限,而 MS 图像则提供丰富的光谱分辨率,但缺乏空间清晰度。通过将来自 MS 和 PAN 图像的互补信息整合为一个复合表示,作者可以实现空间和光谱分辨率之间的有效平衡。类似地,在医学图像融合(MIF)领域,各种成像技术捕获不同类型的信息。例如,计算机断层扫描(CT)图像可以提供关于骨骼和高质量组织的详细洞察,而磁共振成像(MRI)可以提供更高分辨率的图像,具有丰富的软组织细节。凭借各种模式的互补信息,MIF 可以克服单模态图像的局限性,提供更全面和详细的表现,有助于现代医学诊断。

近年来,深度神经网络(DNNs)的繁荣发展导致了多种基于 DNN 的多模态图像融合方法的出现。在超分辨率中,先驱工作 PNN [16] 采用了一个简单的三层神经网络实现显著的结果,这些结果在此之前被认为是不可能的,突显了深度学习的强大学习能力。之后,越来越多的复杂和更深入的架构出现了 [23, 24],展示了出色的视觉性能。然而,现有的多模态融合方法面临共同的局限性。基于卷积神经网络(CNN)的方法常常难以建立全局感受野。尽管通过自注意力机制的 Transformer 解决了这个问题,但也引入了与二次计算复杂度相关的显著挑战。如今,结构状态空间模型已引起了广泛关注,因为它们具有较高的计算效率和用其建模长序列依赖性的原则能力。然而,它们在进行 2D 图像处理时可能引入偏差先验知识。为了克服这些挑战,作者有理由设计一种新颖的顺序扫描方法及其应用框架。

作者的动机:

全球建模能力在图像恢复任务中至关重要,因为图像恢复的一个重要方面就是在图像内部找到有用的信息以弥补当前块中缺失的数据。长期以来,CNN和ViT一直是计算机视觉的主导架构,各自具有独特的优点和局限性。CNN受到局部感受野的限制,阻碍了它们建模长程依赖的能力。相比之下,ViT使用自注意力机制访问全局感受野,但受到二次计算复杂度的限制。最近,结构状态空间模型在捕获序列数据中的长期依赖性时,尽管保持了线性时间复杂度,但已展示了增强的能力。值得注意的是,Mamba通过选择性状态空间和硬件感知算法在降低推理延迟和改进整体性能方面取得了重要进步。随着vmamba [17]和Vision Mamba [16]的引入,人们开始将状态空间模型应用于视觉任务。

然而,目前大多数基于Mamba的方法都采用单向状态空间模型,这些方法在接收视野范围上受到某些限制,因为它们采用扫描方法。具体而言,输入块的初始部分的视野较大,但后一部分的视野变得非常小,这可能会削弱模型的全局建模能力。此外,与语言建模中词序依赖之间的关系不同,图像的二维空间信息呈现给简单的顺序扫描方法带来了相当大的挑战。传统策略,如平铺图像块并顺序扫描它们,可能会引入有偏的局部二维依赖性,从而削弱模型有效识别空间关系的能力。

为解决这些挑战,作者提出了一种新的序列扫描方法,名为随机打乱扫描(Random Shuffle Scanning)。图1展示了三种方法的效验感受野(ERF)的可视化。与卷积方法相比,自注意力方法的优点在于实现了全局感受野。然而,作者的方法提供了整体无偏全局感受野,使得网络能够丢弃固定局部先验并更有效地关注需要学习的内容。对于顺序输入图像块,作者首先应用位置编码,然后对块进行随机打乱,使其经过Mamba块进行长程依赖建模。随机打乱方法在数学期望上消除了局部和全局2D依赖的前置概念,使得模型能够访问无偏先验,并最终建立更一致和全局感受野。

picture.image

作者的贡献可以总结如下:

  1. 作者设计了Shuffle Mamba框架,其中关键组件中的随机打乱操作提供了无偏全局感受野,而不会增加任何参数。
  2. 作者开发了一种针对训练和测试此框架的特定策略。在训练中,每个输入都独立地进行随机打乱操作。在测试中,作者使用蒙特卡罗平均来估计每个Mamba块的输出。

3) 在两个重要的多模图像融合任务上的广泛实验表明,作者的方法在定量评估和视觉质量上均取得了出色的性能。

Related Works

State Space Model

状态空间模型(SSM)源于控制理论,由于其出色的建模能力,在深度学习领域得到了广泛应用,特别是在处理长短期依赖关系。S4模型[13]首次引入了SSM的概念,有效降低了状态表示相关的计算和内存需求,同时实现了全局信息建模。在S4的基础上,S5模型[12]引入了MIMO结构和高效的并行扫描策略,提高了性能,同时并未显著增加计算需求。H3模型[14]进一步优化了这些方法,实现了在语言建模任务中与Transformer类似的性能和效率。

最近,Mamba[13]通过选择性状态空间和硬件感知的算法,在推理速度和性能指标上有了显著提高。Vmamba[17]和Vision Mamba[17]的出现,使得SSM在高性能视觉任务中的应用受到关注。然而,大多数现有的基于SSM的视觉模型[17, 18, 19, 20]采用固定的扫描策略,这可能导致模型在处理低级视觉任务时引入先验偏见。具体来说,固定的图像块选择策略可能导致模型逐渐忽略之前的输入序列,而处理当前块,从而损害其建模全局信息的能力。

为了解决这一挑战,Vmamba[17]引入了CSM,它从前向后、后向前、左向右和右向左等不同方向扫描图像像素。在这想法的基础上,RSM[22]提出了OSSM,将图像块扁平化为八种方向序列,增强了网络捕获和建模大规模空间特征的能力。此外,LocalMamba[17]应用窗口选择扫描,确保全局和局部视觉线索的和谐融合。RS-Mamba[1]则融入了动态多路径激活机制,用于模拟非因果数据。但是,这些方法并没有同时考虑图像结构的完整性和扫描处理过程中的像素随机性,本质上仍然是固定的固定策略。为了解决这个问题,作者推出了Shuffle Mamba框架,它包括随机打乱和逆操作,以获得在全球感受野中没有偏见的结果。

Method

Preliminaries

受连续线性时不变(LTI)系统的启发,自适应小波(SSMs)利用一个隐式潜在状态将一维序列映射到。具体而言,SSMs可以表示为一个常微分方程(ODE):

其中,是演化矩阵,而和充当投影参数。然而,在深度学习的背景下,求解这些微分方程可能会具有挑战性。S4和Mamba模型提出了引入一个时间尺度参数来将连续参数、转换为它们的离散对应物、:

(4) (5)

最后,通过全局卷积可以获得系统的输出:

其中,表示一维序列的长度,是一个结构化卷积核。

Network Architecture

提出的Shuffle Mamba框架包含三个功能组件:随机Mamba块(RM块)、随机通道交互Mamba块(RCIM块)和随机模态交互Mamba块(RMIM块)。整体工作流程如图2所示。假设输入图像具有不同模态,分别为和,其中作为上采样,作者首先使用卷积层将图像投影到特征空间。由于卷积层的感受野有限,捕捉全局特征具有挑战性,因此作者进行填充嵌入,将生成的填充嵌入块送入RM块进行全局特征提取。这一过程产生了全局模态特定的特征和,其中:

picture.image

全局模态特定的特征随后被发送到RCIM块进行简单的通道信息交换,无需引入额外的参数。经过各自RM块处理交换后的特征得到和。接下来,作者使用RMIM块通过模态特征和的深度融合获得。从而,最终的融合图像可以通过reshape和对角卷积连接进行访问:

Key Components

随机摇摆扫描。Mamba最初是为了适应语言序列的建模而被设计的。为了访问在2D图像中保持全局感受野和局部依赖性的无偏独立性,作者提出了随机摇摆扫描方法。如图3所示,对于连续输入2D图像块,作者首先对位置建模应用深度卷积。图像块然后被随机摇摆,并被送到Mamba块进行长程依赖性建模。这种策略使Mamba块有等概率模拟相邻块之间的相互作用,从而使网络有效地学习和建模一个无偏的先验。此外,因为块的相对位置对于恢复语义信息至关重要,输出图像必须根据逆摇摆精确地与输入对齐。因此,随机摇摆及其逆操作组成了一个信息损耗为零的转换对。

picture.image

随机Mamba块。基于这个摇摆-逆派对的随机Mamba块的设计,首先对输入特征 进行层规范化,得到 ,然后通过随机摇摆和多层感知器(MLPs)将其投影到 和 。在第一个分支中, 通过带有SiLU激活的1D卷积层,生成 。SSM(单流多模态)用于计算输出 。而在另一个分支中, 被发送到激活函数以生成 的开门参数。最后,作者应用逆摇摆和残差连接得到最终输出序列 。

随机通道交互Mamba块。在RCIM块中,作者从[1]中采用了实现不同模ality之间轻量级特征交互的方法。作者使用分裂操作将模ality特征 和 基于通道维度分成两半,然后进行互补拼接。交换的特征然后被送至各自的RM Blocks进行处理。通过重复这些步骤,全局的模ality-specific特征被 initialized。

随机模式交互Mamba块。受到交叉注意力启发,作者设计了RMIM块以处理多模态图像信息。在这种方法中,作者将随机摇摆序列特征映射到共享空间,并使用 Mask 机制在无偏先验下学习互补信息,从而减少冗余特征对融合结果的影响。作者采用与RM块类似的方法生成 和 ,并使用输入 生成 Mask 参数 进行动态调整。两个输出被组合并投影,然后通过逆摇摆和形状调整操作与输入序列对齐。最后,模块的输出 通过深度卷积和特征 flattening 得到。

Testing with Monte Carlo averaging

作者将随机因素引入随机混洗操作,这需要在这些因素在生成最终融合时进行边缘化。然而,随机混洗方法在理论上是具有一定挑战性的,因为存在指数大小的潜在模型,使得对这些模型的预测进行精确平均变得不可能。作者从[14]中的dropout中得到启发,通过分层期望来近似整个模型的期望值。因此,在测试期间随机混洗的计算可以表示如下,其中:

实际上,根据上述方程估计需要遍历所有可能的混洗结果,这给计算带来很大的负担。因此,作者使用蒙特卡洛平均来估计其期望:

具体来说,输入图像独立地混洗了次,然后计算了的个输出。这些输出的平均值用于获得最终估计。当时,蒙特卡洛估计器近似地逼近真实平均值。图4说明了作者设计的测试过程,通过将多个相同输入放入子批次并在GPU上并行计算,极大地减少了实际测试时间。

picture.image

4 Loss Function

如图5:几种方法在WV3数据集上的视觉对比实验

picture.image

根据该领域的惯例,作者模型在去模糊化的损失函数选择L1损失。在MIF任务中,作者使用两个输入图像和融合图像来计算一个复合损失函数,包括L1损失、SSIM损失和梯度损失。## 实验

Datasets and Benchmark

针对全重构任务,作者使用了WorldView-II(WV2),高芬2(GF2)和WorldView-III(WV3)的数据集,包括各种城市和自然场景。作者在不可见 GT 值中遵循Wald [21]协议生成训练样本。作者对提出的方法与经典方法和技术进行了全面比较,包括INNformer [22],SFINet [23],MSDDN [17],PanFlowNet [24],FAME [16],DISPNet [25]和Pan-mamba [15]。

图6:在MRI-PET数据集上,几种方法进行了比较视觉效果实验multi-modal融合技术,包括PSLPT [22],EMFusion [23],MSRPAN [24],SwinFusion [25],Zero [1],U2Fusion [23],和CDDFuse [22]。

picture.image

在MIF任务中,作者使用了Harvard Medical网站上的一系列医学图像,包括MRI-CT,MRI-PET和MRI-SPECT的图像对。在这个任务中,作者将作者的方法与各种基于深度学习的方法进行了比较。

所有实验都是在PyTorch框架上使用两个NVIDIA RTX 3060 GPU进行的。作者在全重构任务(pan-sharpening task)上训练了500个周期,批次大小为4,在MIF任务(medical image fusion(MIF)任务)上训练了200个周期,批次大小为1。使用Adam优化器优化网络参数。初始学习率设置为,在指定的周期中使用CosineAnnealingLR调度器将其减少到。对于全重构任务,作者将训练集图像随机裁剪为32x32的小块和128x128的PAN图像。对于MIF任务,将训练集图像裁剪到256x256。

Comparison with SOTA Methods

Pan-sharpening.

表1中列出了在三个数据集上的实验结果。参考指标,包括峰值信噪比(PSNR)、结构相似性指数(SSIM)、均方误差(SAM)和ERGAS[1],用于评估融合效果。结果表明,所提出的方法在所有指标上都优于现有的最佳方法。尤其是,在PSNR指标上,作者的方法比Pan-mamba[2,3]实现了0.1047、0.2727和0.2301 dB的提升,位列第二。更好的PSNR和SSIM证明了融合结果从原始图像中传递了更多的信息且对 distortion 的影响更小。在定性对比方面,图5展示了来自WV3数据集的代表性样本。作者的方法在MSE图中表现更好,说明融合结果与 GT 值更接近。与其他方法相比,作者的方法在恢复光谱和空间细节方面更加准确,突出了作者融合技术的优势。

picture.image

为了进一步验证作者的方法在全分辨率场景下的泛化能力,作者使用三个非参考指标,包括 、 和 QNR,在完整的WorldView-II数据集上评估作者的方法。表2呈现了实验结果。作者的方法在所有结果上都优于对比方法,展示了 Shuffle Mamba 在图像融合方面的强健适应性。

picture.image

Medical Image Fusion.

表3显示了在MIF数据集上四个指标的定量对比结果。所提出的方法在各个指标上表现良好,显示了其在医学图像融合方面的有效性。在我国的实验结果中,更高的VIF表明与人类感知更接近。改进的SCD、Qabf和SSIM分数表明融合图像保持了更高的相似性,并相对于原始图像经历了更少的失真。图6在MRI-PET数据集上展示了几种方法的定性比较。作者的方法在视觉质量上表现出优越性,这一结论得到了实验指标的支持。

picture.image

Ablation Experiments

为了验证随机重排操作的有效性,作者进行了相应的消融实验。这些实验的核心操作包括从三个主要模块中移除重排操作。值得注意的是,在移除与RM相关的组件的实验中,排除了RCIM模块中的RM块。实验结果如表4所示。移除重排操作导致模型性能显著下降,表明在Shuffle Mamba框架下设计的新扫描策略有效地提高了多模态图像融合的质量。

picture.image

由于Monte Carlo平均,Shuffle Mamba使得在提高性能的同时利用内存和计算资源成为可能。作者研究了样本数量、PSNR指数和资源消耗之间的关系。作者对每个样本大小进行了五次实验,计算各种指标的平均和标准差。图7(a)说明了PSNR的趋势,图7(b)和(c)详细介绍了相应的内存使用和处理每个图像所需的时间。随着样本数量的增加,性能和资源消耗上升,可以实现性能和效率之间的权衡。此外,Monte Carlo平均增强了随机重排操作的理论鲁棒性,有效提高了PSNR的平均值,同时减小了融合结果的方差。更多的实验结果可在附录中找到。

picture.image

Conclusion

在本文中,作者用一种基于随机Shuffle的方法替换了当前Mamba方法中使用的固定扫描策略,并设计了一个新的Shuffle Mamba框架。这种方法可以减轻固定扫描方法带来的偏见,并为多模态图像融合提供全局接收范围。

在测试中,作者采用蒙特卡洛平均值来考虑引入的随机因素。在两项任务上的大量实验表明,作者的方法超过了最先进的方法,并展示了强大的泛化能力。

参考

[1].Shuffle Mamba: State Space Models with Random Shuffle.

点击上方卡片,关注 「AI视界引擎」 公众号

0
0
0
0
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论