ECCV2024｜性能超越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中的策略——由输入生成卷积权重（类似于动态卷积，所以表达能力还是很强的）
• 使用四个方向的扫描，让当前像素的输出实际上是来自上下左右四个方向邻域的一维卷积结果

2. 既然可以做，那么效果怎么样呢？

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

动机

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参数来调整来自残差连接中的权重：