ECCV 2024 | ModelMambaIR：基于Mamba模型用于图像恢复的简单基线

小白学视觉

发布于 2024-12-23 06:27:15

3570

文章被收录于专栏：深度学习和计算机视觉深度学习和计算机视觉

论文信息

题目：MambaIR: A Simple Baseline for Image Restoration with State-Space Model

MambaIR：基于Mamba模型的图像恢复简单基线

作者：Hang Guo, Jinmin Li, Tao Dai, Zhihao Ouyang, Xudong Ren, Shu-Tao Xia

源码：https://github.com/csguoh/MambaIR

论文创新点

首次将状态空间模型应用于图像恢复：作者首次将Mamba这一先进的状态空间模型引入到图像恢复任务中，提出了MambaIR，作为CNN和Transformer方法的简单但有效的替代方案。
提出残差状态空间块（RSSB）：为了解决标准Mamba在图像恢复中的局部像素遗忘和通道冗余问题，作者设计了残差状态空间块（RSSB），通过引入局部卷积和通道注意力机制，显著提升了模型的性能。
全局感受野与线性复杂度的结合：MambaIR通过并行扫描算法和状态空间方程，实现了全局感受野与线性复杂度的结合，解决了现有方法在全局建模和高效计算之间的权衡问题。
广泛的实验验证：作者在多个图像恢复任务（如图像超分辨率、图像去噪等）上进行了广泛的实验验证，结果表明MambaIR在多个基准数据集上均优于现有的强基线方法，展示了其在图像恢复中的强大潜力。

摘要

近年来，图像恢复领域取得了显著进展，主要归功于现代深度神经网络（如CNN和Transformer）的发展。然而，现有的恢复骨干网络在全局感受野和高效计算之间面临困境，限制了其在实际应用中的应用。最近，选择性结构化状态空间模型（尤其是改进版Mamba）在长程依赖建模方面展现出巨大潜力，具有线性复杂度，为解决上述困境提供了一种方法。然而，标准的Mamba在低级视觉任务中仍面临局部像素遗忘和通道冗余等挑战。在本文中，作者提出了一种简单但有效的基线模型，名为MambaIR，通过引入局部增强和通道注意力来改进原始Mamba。通过这种方式，MambaIR利用了局部像素相似性并减少了通道冗余。大量实验证明了作者方法的优越性，例如，MambaIR在图像超分辨率任务中比SwinIR高出0.45dB，且具有相似的计算成本但拥有全局感受野。

关键词

图像恢复，状态空间模型， Mamba

3 方法

总体架构

如图2所示，MambaIR由三个阶段组成：浅层特征提取、深层特征提取和高质量重建。给定低质量（LQ）输入图像，作者首先使用卷积层从浅层特征提取中生成浅层特征，其中和表示输入图像的高度和宽度，是通道数。随后，浅层特征经过深层特征提取阶段以获取第层的深层特征，。该阶段由多个残差状态空间组（RSSG）堆叠而成，每个RSSG包含若干残差状态空间块（RSSB）。此外，在每个组末尾引入额外的卷积层以细化从RSSB提取的特征。最后，作者使用元素级求和来获得高质量重建阶段的输入，用于重建高质量（HQ）输出图像。

残差状态空间块

在之前的基于Transformer的恢复网络中，块设计主要遵循Norm Attention Norm MLP流程。尽管Attention和SSM都可以建模全局依赖性，但作者发现这两个模块的行为不同（详见补充材料），简单地将Attention替换为SSM只能获得次优结果。因此，为基于Mamba的恢复网络定制全新的块结构是有前景的。

为此，作者提出了残差状态空间块（RSSB）以适应SSM块用于恢复。如图2(a)所示，给定输入深层特征，作者首先使用LayerNorm（LN），然后使用视觉状态空间模块（VSSM）捕捉空间长期依赖性。此外，作者还使用可学习比例因子来控制跳跃连接的信息：

此外，由于SSM处理展平特征图作为1D token序列，序列中邻近像素的数量受到展平策略的极大影响。例如，当采用四向展开策略时，锚点像素只能感知四个最近的邻居（见图3(a)），即2D特征图中的某些空间接近像素在1D token序列中距离较远，这种过度距离可能导致局部像素遗忘。为此，作者在VSSM后引入额外的局部卷积以帮助恢复邻近相似性。具体来说，作者首先对进行LayerNorm归一化，然后使用卷积层补偿局部特征。为了保持效率，卷积层采用瓶颈结构，即通道首先通过因子压缩以获得形状为的特征，然后进行通道扩展以恢复原始形状。

此外，SSM通常引入较大的隐藏状态数量以记忆非常长程的依赖关系，作者在图3(b)中可视化了不同通道的激活结果，发现存在显著的通道冗余。为了促进不同通道的表达能力，作者在RSSB中引入了通道注意力（CA）。通过这种方式，SSM可以在后续通道注意力选择关键通道后，专注于学习多样化的通道表示，从而避免通道冗余。最后，另一个可调比例因子用于残差连接以获得RSSB的最终输出。上述过程可以表示为：

视觉状态空间模块

为了保持效率，基于Transformer的恢复网络通常将输入划分为小patch或采用移位窗口注意力，阻碍了整个图像级别的交互。受Mamba在长程建模中具有线性复杂度的成功启发，作者引入了视觉状态空间模块用于图像恢复。

视觉状态空间模块（VSSM）可以通过状态空间方程捕捉长程依赖性，VSSM的架构如图2(b)所示。遵循[44]，输入特征将通过两个并行分支。在第一个分支中，特征通道通过线性层扩展到，其中是预定义的通道扩展因子，随后是深度卷积、SiLU激活函数、2D-SSM层和LayerNorm。在第二个分支中，特征通道也通过线性层扩展到，随后是SiLU激活函数。之后，两个分支的特征通过Hadamard积进行聚合。最后，通道数投影回以生成与输入形状相同的输出：

其中DWConv表示深度卷积，表示Hadamard积。

2D选择性扫描模块

标准的Mamba以因果方式处理输入数据，因此只能捕捉扫描部分数据内的信息。这种特性非常适合涉及顺序性质的NLP任务，但在转移到非因果数据（如图像）时面临显著挑战。为了更好地利用2D空间信息，作者遵循[44]并引入了2D选择性扫描模块（2D-SSM）。如图2(c)所示，2D图像特征通过沿四个不同方向（左上到右下、右下到左上、右上到左下、左下到右上）扫描展平为1D序列。然后根据离散状态空间方程捕捉每个序列的长程依赖性。最后，所有序列通过求和操作合并，随后进行重塑操作以恢复2D结构。