超越SwinIR，Mamba入局图像复原，达成新SOTA

https://arxiv.org/abs/2402.15648 https://github.com/csguoh/MambaIR

本文概述

最近，选择性结构化状态空间模型（例如 Mamba）在具有线性复杂性的远程依赖关系建模方面表现出了巨大的潜力，但它在低级计算机视觉中仍处于探索之中。

在这项工作中，我们引入了一个简单但强大的基准模型，名为 MambaIR，用于图像恢复。具体来说，我们提出残差状态空间块作为核心组件，它采用卷积和通道注意力来增强普通曼巴的能力。通过这种方式，我们的 MambaIR 利用局部补丁重现先验以及通道交互来生成特定于恢复的特征表示。

大量实验证明了我们方法的优越性，例如，使用类似的计算成本但具有全局感受野，MambaIR 的性能比基于 Transformer 的基线 SwinIR 提高了 0.36dB。

本文贡献

• 第一个通过大量实验制定 MambaIR 来调整状态空间模型以进行低级图像恢复，MambaIR 是 CNN 和 Transformer 的简单但有效的替代方案。
• 提出了剩余状态空间块（RSSB），它可以通过局部空间先验和通道交互来增强原始Mamba 的能力。
• 对各种任务的广泛实验表明，所提MambaIR 优于基于 Transformer 的强大基线，可为图像恢复提供强大且有前途的骨干模型解决方案。

本文动机

动机部分引自作者解析《性能超越SwinIR！MambaIR: 基于Mamba的图像复原基准模型》，链接为https://zhuanlan.zhihu.com/p/684248751

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模型，并方便之后的进一步研究改进。因此在模型设计上，作者遵循了之前图像复原任务的大致流程，即浅层特征提取，深层特征提取，以及高质量图像重建。

我们提出了剩余状态空间块（RSSB）来改进具有特定恢复先验的原始SSM块。具体来说，我们考虑局部补丁重复以及通道交互以协助Mamba中的远程空间成型。如图2（a）所示，给定输入深度特征

我们首先使用Layernorm（LN），然后使用视觉状态空间模块（VSSM来捕获空间长-期限依赖性。之后，我们还在跳过连接中使用可学习的比例因子s以获得更好性能：

之后，我们使用另一个 LayerNorm 来规范化

.然后使用卷积层对空间局部相似性先验进行建模。为了减轻通道冗余和效率，卷积层采用瓶颈结构，即通道首先按因子

压缩得到形状

的特征，然后我们进行通道扩展以恢复原始形状。此外，为了对通道之间的交互进行建模，我们引入了通道注意层（CA）[22]，然后在残差连接中引入另一个可调比例因子以获得最终输出。该过程描述如下：

Vision State-Space Module

为了保持效率，基于 Transformer 的恢复网络通常将输入分成小块 或采用移位窗口注意力，从而阻碍了整个图像级别的交互。受到 Mamba 在具有线性复杂性的远程建模方面的成功的激励，我们将视觉状态空间模块引入图像恢复。计算过程如下：

本文实验