MatIR：融合Mamba与 Transformer 的混合图像修复模型，实验验证其显著有效性！

近年来，基于Transformer的模型通过利用其内在捕捉复杂上下文特征的能力，在图像修复领域取得了显著进展。近期，Mamba模型凭借其处理长程依赖的能力和相较于Transformer的显著计算效率，在计算机视觉领域引起了广泛关注。 然而，Mamba在上下文学习能力方面目前落后于Transformer。 为了克服这两种模型的局限性，作者提出了一种名为MatIR的Mamba-Transformer混合图像修复模型。 具体来说，MatIR交叉循环Transformer层和Mamba层的块以提取特征，从而充分利用这两种架构的优势。 在Mamba模块中，作者引入了图像修复状态空间（IRsS）模块，它沿着四个扫描路径遍历，以实现长序列数据的有效处理。 在Transformer模块中，作者将基于三角窗口的局部注意力与基于通道的全局注意力相结合，以在更广泛的图像像素范围内有效地激活注意力机制。大量的实验结果和消融研究表明了该方法的有效性。

引言

图像恢复旨在从退化或损坏的输入中恢复清晰且高质量的图像。这是计算机视觉中的一个长期问题，包括了一系列子问题，如超分辨率、图像去噪和去模糊。随着现代深度学习模型，如卷积神经网络CNNs 和 Transformer 的引入，近几年来，该领域的最先进性能得到了持续提升。去噪、去模糊和超分辨率等任务需要能够准确重建图像细节同时保留结构信息的模型。传统的基于卷积的方法通常无法捕捉到涉及大或严重退化图像的任务中至关重要的长距离依赖关系。深度学习领域的最新进展，如Transformer架构，在捕捉图像中的全局依赖性方面显示出了潜力。然而，Transformer的计算成本随序列长度的平方增长，这限制了其可扩展性，尤其是在高分辨率图像恢复任务中，这些任务以二次复杂度为代价提供了全局感受野。

近期，由于其在处理长距离依赖关系和相对于Transformer在计算效率上的显著优势，Mamba在计算机视觉领域日益突出。Mamba架构是一种新的序列模型，通过引入状态空间模型（SSM）的概念，实现了对长序列数据的有效处理。状态空间模型（SSM）利用状态空间表示来达到线性计算复杂度，并且能够在不牺牲准确性的情况下高效地处理长序列。与序列长度相关的线性计算复杂度显示了其在序列建模任务中处理长距离依赖关系的高效性和有效性。然而，现有研究显示，在上下文学习（ICL）能力方面，Mamba落后于Transformer。[14, 19, 50] 从现代深度神经网络发展的最新进展中汲取灵感。在本研究中，作者的目标是弥补这些模型的不足，同时利用它们的各自优势。作者提出了一种混合Mamba-Transformer图像恢复模型，称为MatIR。这是一个新颖的混合架构，它结合了Mamba架构的优势（在处理长序列时以内存效率著称）和Transformer（在上下文学习和信息检索方面表现卓越）[14, 19, 50]。通过结合这两种方法，MatIR旨在为各种图像恢复任务提供一个强大而高效的解决方案。[33, 45]

具体而言，1）浅层特征提取阶段采用简单的卷积层提取浅层特征。随后，2）深度学习特征提取阶段使用堆叠了Mamba层的Transformer。在Transformer层中，作者运用三角窗局部注意力（TWLA）和通道全局注意力（CGA）机制，有效激活了更广泛的图像像素范围，从而提升了该模块的性能。在Mamba层中，作者运用图像恢复状态空间（IRSS）模块，通过沿四个扫描路径进行遍历，实现对不同方向和路径的长序列数据的有效处理。在提升性能和吞吐量的同时，保持可管理的内存占用。作为MatIR的三个核心组件，TWLA和CGA分别激活了局部和全局范围内，三角窗和矩形窗中的更多输入像素，实现了更高品质的图像恢复。IRSS通过四个不同方向路径（左、上、右、下）在状态空间中创建与序列长度相关的线性计算复杂度信息，展示了在序列建模任务中处理长距离依赖的效率和有效性。最后，3）高质量图像重建阶段聚合浅层和深层特征，生成高质量的输出图像。凭借局部和全局有效的感受野、高效的内存管理和计算，MatIR成为图像恢复主干网络的新选择。

总的来说，作者的主要贡献可以概括如下：

1. 作者通过广泛实验将状态空间模型和Transformer注意力机制应用于图像修复领域，从而制定了MatIR，该算法在保持性能的同时大幅提升了计算效率。作者提出了图像修复状态空间（IRSS）模块，该模块沿着四个不同的路径进行扫描和遍历，以实现对长序列数据的高效处理。
2. 作者提出了一种三角形窗口局部注意力（TWLA）模块和一种通道全局注意力（CGA）模块，分别用于从局部和全局范围内、三角形窗口和矩形窗口中激活更多输入像素，以实现更高质量的图像修复。
3. 通过在多个基准数据集上的全面评估，作者的方法在性能上优于其他最先进的方法，为图像恢复提供了一个强大且前景广阔的 Backbone 解决方案。

2. 相关研究工作

图像恢复。图像恢复是计算机视觉领域的一个长期问题。在过去十年中，在包括图像恢复在内的多个领域，人们投入了大量努力以提升深度学习方法的性能。由SRCNN [16]开创，深度学习通过一个简单的三层卷积神经网络（CNN）被引入到图像恢复超分辨率中。自那以后，许多研究探索了各种架构改进以提升性能。VDSR [23]实现了一个更深层的网络，而DRCN [24]提出了递归结构。EDSR [31]和RDN [69]开发了新的残差模块，以进一步改善CNN在超分辨率中的能力。然而，尽管CNN取得了成功，但其感受野本质上有限，这使得捕捉长距离依赖关系变得困难。

近年来，视觉Transformer（ViT）[17]及其变体[11, 34, 52]将自注意力机制引入图像处理领域，使模型能够学习全局关系。基于此，IPT[6]成功尝试利用基于Transformer的网络进行各种图像恢复任务。此后，开发了多种技术来提升图像恢复Transformer的性能。这些技术包括SwinIR[30]和CAT[9]实现的移位窗口自注意力、ELAN[66]的分组多尺度自注意力机制、ART[59]和OmniSR[51]的稀疏注意力，以及GRL[29]的 Anchor 定自注意力机制、多注意力机制DART和DISR[54, 55]，它们都旨在扩大感受野范围，以获得更好的效果。然而，这以二次计算复杂度为代价提供了全局感受野。自注意力在序列长度方面的二次计算复杂性问题构成了挑战，尤其是在处理高分辨率图像时。

状态空间模型。近年来，由于能够处理长距离依赖关系以及相对于Transformer的高计算效率，Mamba在计算机视觉领域日益突出。Mamba架构是一种新的序列模型，通过引入状态空间模型（SSM）的概念[14, 19, 50]，实现了对长序列数据的有效处理。状态空间模型（SSM）利用状态空间表示，实现线性计算复杂度，并能高效处理长序列数据而不影响准确性。与序列长度相关的线性计算复杂度显示出在序列建模任务中处理长距离依赖关系的高效性和有效性。然而，现有研究表明，Mamba在上下文学习（ICL）能力上落后于Transformer[50]。在高效计算和全局建模之间的权衡困境尚未得到根本解决[20, 33, 45]。考虑到两种当前最先进模型的局限性，作者探讨了混合Mamba-Transformer方法在图像修复中的潜力。在本论文中，基于Mamba和Transformer模型的有效性，作者提出了三个核心组件：图像修复状态空间（IRSS）模块，该模块沿着四种不同的路径扫描和遍历，以实现长序列数据的有效处理。三角窗口局部注意力（TWLA）块和通道全局注意力（CGA）块，每个块都通过注意力机制激活来自局部和全局、三角形和矩形窗口范围内更多的输入像素，以实现更高品质的图像修复。

3. 研究方法

3.1. 前言

3.2 全局注意力模块（Channel Global AttentionBlock）

3.3 三角形窗口局部注意力块

3.4 图像恢复状态空间块

Transformer注意力机制的计算成本随着序列长度的增加呈二次增长，这限制了它们的可扩展性，尤其是在高分辨率图像恢复任务中。状态空间模型（SSM）利用状态空间表示实现线性计算复杂度，可以高效处理长序列而不影响准确性。受这些特性的启发，作者在恢复模型设计中引入了图像恢复状态空间模块（见3.1及图3），为了最大化该模块的性能，IRSS从四个不同的方向路径（左、上、右、下）在状态空间中创建了与序列长度相关的线性计算复杂度信息。设计的IRSS模块作为MatIR模型的独立Mamba层，与Transformer层进行交互。

图像恢复状态空间模块块（Image Restoration State Space Module Block，简称IRSS）

3.5. MatIR总体架构

4. 实验

4.1. 实验设置。

数据集与评估。在借鉴前人工作[30, 59]的基础上，作者对包括图像超分辨率和图像去噪（即高斯彩色图像去噪和现实场景去噪）以及散焦去模糊效果（即S：单图像散焦去模糊。D：双像素散焦去模糊）在内的各种图像恢复任务进行了实验。作者采用DIV2K[49]和Flickr2K[32]来训练经典超分辨率模型。此外，作者使用Set5[5]、Set14[58]、B100[37]、Urban100[21]和Manga109[38]来评估不同超分辨率方法的有效性。对于高斯彩色图像去噪，作者利用DIV2K[49]、Flickr2K[32]、BSD500[4]和WED[36]作为作者的训练数据集。高斯彩色图像去噪的测试数据集包括BSD68[37]、Kodak24[18]、McMaster[64]和Urban100[21]。对于现实场景去噪，作者使用SIDD[1]数据集的320张高分辨率图像来训练作者的模型，并使用SIDD测试集和DND[44]数据集进行测试。遵循[30, 70]，当测试中使用自集成策略[32]时，作者将该模型称为MatIR。性能评估采用PSNR和SSIM在YCbCr颜色空间的Y通道上进行。散焦去模糊：作者在包括户外、室内和混合场景的数据集[2,22,27,28,48,53,57]上与12种最先进的方法进行了比较。

训练细节。根据先前的研究[8, 30, 59]，作者在训练过程中采用水平翻转和随机旋转、和的数据增强方法。此外，作者将原始图像裁剪成的块用于图像超分辨率，以及的块用于图像去噪。在图像超分辨率方面，作者使用模型的预训练权重来初始化和的权重，并将学习率及总训练迭代次数减半以缩短训练时间[32]。为确保公平比较，作者将图像超分辨率的训练批次大小调整为32，图像去噪调整为16。作者采用Adam[25]作为优化器，MatIR的，。初始学习率设为，当训练迭代达到特定里程碑时将学习率减半。作者的MatIR模型在8块NVIDIA A100 GPU上训练。

评估指标。在大多数图像恢复任务中，作者采用峰值信噪比（PSNR）、结构相似性指数（SSIM）和LPIPS作为评估指标。一般来说，更高的PSNR和SSIM值，以及更低的LPIPS和FID值，意味着更好的性能。

4.2. 消融研究

不同Transformer层设计的影响。作为核心组件，TWLA和CGA通过激活更多输入像素，实现局部和全局的三角形注意力窗口和矩形注意力窗口范围，从而为MatIR提供更高品质的图像恢复。在本节中，作者对这两个关键组件分别进行了消融研究。如表1和表2所示的结果：（1）在Transformer层中将这两个组件 ConCat 使用比单独使用每个组件具有更大的增益效果；（2）如果Mamba层的核心组件IRSS采用Transformer方法，计算需求将增加，但收益相同。不同扫描模式在IRSS中的影响。为了使Mamba处理2D图像，需要先将特征图展平，然后通过状态空间方程迭代。因此，展开策略尤为重要。在本研究中，作者遵循[33]的方法，使用四种不同的扫描方向生成扫描序列。在此，作者消融了不同的扫描模式以研究其影响，与单方向（左上到右下）和双向（左上到右下，右下到左上）扫描相比，使用四向扫描允许 Anchor 点像素感知更大的邻域范围，从而获得更好的结果。

4.3 图像超分辨率评估

作者将在5个公开的经典超分辨率数据集上，将作者的方法与16种最先进的检索方法进行比较。定量结果展示在表4中。作者可以看到，作者的方法在5个不同的数据集上优于大多数方法。特别是，与SRformer方法相比，作者的Urban100 X4方法在PSNR上领先SRformer高达0.87，在SSIM上领先DDNM高达0.0194dB。对于定性结果，作者的方法具有最佳的视觉效果，包括更逼真的纹理，如图5所示。这些视觉比较与定量结果一致，证明了作者方法的有效性。更多视觉结果详见补充材料。

4.4 图像去噪评估

表5展示了高斯颜色图像去噪的结果。与[60, 63]中类似，比较的噪声水平包括15、25和50。可以看出，作者的模型在大多数数据集上实现了最佳性能。特别是，它在Urban100数据集上超过了SwinIR[30]，甚至达到了0.51dB（σ=50）。作者还在图5中给出了视觉比较。得益于全局感受野，作者的MatIR能够实现更好的结构保留，从而产生更清晰的边缘和自然的形状。

表6. 失焦去模糊结果。S：单图失焦去模糊。D：双像素失焦去模糊。

4.5 图像去模糊评估

在图像去模糊方面，作者主要评估了失焦去模糊的效果，包括单图像失焦去模糊和双像素失焦去模糊。在表5中，定量结果显示，作者的方法在所有数据集上均取得了最佳性能。与Restormer [57]在户外场景中的表现相比，作者方法的PSRN提升可高达0.73dB。作者的图像具有最佳的视觉质量，细节更加真实，接近真实图像。由于篇幅限制，作者在补充材料中提供了更多定量和定性的结果。

4.6. 真实图像去噪

作者进一步转向真实图像去噪任务，以评估MatIR在面临现实世界退化时的鲁棒性。遵循[57]中的方法，作者采用渐进式训练策略以确保公平比较。结果如表7所示，作者的方法在现有最先进模型中实现了最佳性能，并在SIDD数据集上比其他方法（如Uformer[53]）高出了0.31 dB的PSNR，展示了作者方法在真实图像去噪方面的能力。

4.7. 进一步实验

运行时间。作者将所提出方法的运行时间与SOTA信息检索（IR）方法进行了比较。为了进行公平的比较，作者使用各方法公开可用的代码，在NVIDIA TITAN RTX上对输入图像的所有方法进行了评估。如表3所示，所提出方法的运行时间显著优于其他基于Transformer的方法，展示了作者模型良好的计算效率。

参考

[1]. MatIR: A Hybrid Mamba-Transformer Image Restoration Model .