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

最新推荐文章于 2024-05-11 17:16:33 发布

AIWalker-Happy

最新推荐文章于 2024-05-11 17:16:33 发布

阅读量1.9k

点赞数 40

分类专栏：图像超分图像复原 Mamba 文章标签：图像超分 Mamba 图像降噪

本文链接：https://blog.csdn.net/huohu728/article/details/136360575

版权

图像超分同时被 3 个专栏收录

13 篇文章 2 订阅

订阅专栏

图像复原

7 篇文章 2 订阅

订阅专栏

Mamba

1 篇文章 0 订阅

订阅专栏

本文首发:AIWalker

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）所示，给定输入深度特征 $F_D^l \in \mathbb{R}^{H \times W \times C}$ 我们首先使用Layernorm（LN），然后使用视觉状态空间模块（VSSM来捕获空间长-期限依赖性。之后，我们还在跳过连接中使用可学习的比例因子s以获得更好性能：
$Z^l = VSSM(LN(F_D^l)) + s\cdot F_D^l$

之后，我们使用另一个 LayerNorm 来规范化 $Z^l$ .然后使用卷积层对空间局部相似性先验进行建模。为了减轻通道冗余和效率，卷积层采用瓶颈结构，即通道首先按因子 $\gamma$ 压缩得到形状 $R^{H\times W \times \frac{C}{\gamma}}$ 的特征，然后我们进行通道扩展以恢复原始形状。此外，为了对通道之间的交互进行建模，我们引入了通道注意层（CA）[22]，然后在残差连接中引入另一个可调比例因子以获得最终输出。该过程描述如下：
$F_D^{l+1} = CA(Conv(LN(Z^l)))+ s^{'} \cdot Z^l$