【CVPR2025】MambaIRv2: Attentive State Space Restoration-CSDN博客

MambaIRv2: Attentive State Space Restoration

论文信息

题目： MambaIRv2: Attentive State Space Restoration
MambaIRv2：注意力状态空间恢复
源码：https://github.com/csguoh/MambaIR

论文创新点

核心组件创新
- 注意力状态空间方程（ASE）：在Mamba原始状态空间方程基础上，融入提示学习机制。通过构建提示池并设计特定路由策略，让模型能在扫描序列外查询语义相似像素，突破Mamba因果性限制。同时，该方程实现单方向扫描，避免多方向扫描的高计算成本与冗余，提升效率。
- 语义引导邻域（SGN）：鉴于Mamba因果建模导致的长距离衰减问题，SGN根据像素语义标签重构图像，使语义相似像素在一维序列中空间距离拉近。这样能减轻长距离衰减影响，促进远距离相似像素间的有效交互，增强模型对相关像素的利用能力。
整体性能提升
- 高效全局建模：将ASE和SGN等模块集成，赋予Mamba类似ViT的非因果建模能力，实现注意力状态空间恢复。在多种图像恢复任务中，相比同类方法，MambaIRv2在参数更少的情况下，仍能取得更高的峰值信噪比（PSNR），如在轻量级超分辨率任务中比SRFormer的PSNR高出0.35dB ，展示出卓越的性能与效率平衡。
- 广泛任务适应性：在图像超分辨率、JPEG压缩伪影减少、高斯彩色图像去噪等任务上均表现出色，在各任务的多个数据集上达到最优性能，证明其作为通用图像恢复骨干网络的有效性。

摘要

基于Mamba的图像恢复骨干网络最近在平衡全局信息接收和计算效率方面展现出巨大潜力。然而，Mamba固有的因果建模局限性，即每个标记仅依赖于扫描序列中的前序标记，限制了图像中像素的充分利用，从而给图像恢复带来了新的挑战。在这项工作中，我们提出了MambaIRv2，它赋予Mamba类似于视觉Transformer（ViTs）的非因果建模能力，从而实现注意力状态空间恢复模型。具体而言，所提出的注意力状态空间方程允许在扫描序列之外进行注意力聚焦，并仅通过一次扫描即可实现图像展开。此外，我们进一步引入了语义引导邻域机制，以促进远距离但相似像素之间的交互。大量实验表明，我们的MambaIRv2在轻量级超分辨率任务中，即使参数减少了9.3%，峰值信噪比（PSNR）也比SRFormer高出0.35dB，在经典超分辨率任务中，对HAT的PSNR提升高达0.29dB。代码可在https://github.com/csguoh/MambaIR上获取。

1. 引言

1.1 图像恢复研究现状

图像恢复旨在从低质量观测中恢复高质量图像，解决诸如图像超分辨率、图像去噪、JPEG压缩伪影减少等各种子问题。随着深度学习的出现，人们不断取得最先进的性能。早期工作主要利用卷积神经网络（CNNs）作为骨干网络[13, 15, 34, 56, 65]。后来，视觉Transformer（ViTs）[17]因其卓越性能而受到欢迎[6, 8, 11, 32, 33]。最近，选择性状态空间模型（Mamba）[21]得到了探索，显示出作为图像恢复任务替代骨干网络的巨大潜力[23, 45]。

1.2 现有问题及解决方法

尽管基于Mamba的方法具有潜力，但它们面临着重大挑战，特别是由于其对因果状态空间建模的依赖。具体而言，现有方法[23]通过预定义的扫描规则展开二维图像以生成一维标记序列。然而，在Mamba中，每个像素仅基于扫描序列中的前序像素进行建模，即因果属性，这给非因果图像恢复任务带来了一些不利影响。首先，如图1（a）所示，查询像素只能从其前序像素捕获信息，无法感知后续像素，这导致图像中有用像素的利用不足。其次，固有的因果属性导致需要进行多方向扫描，现有方法[23, 45, 49]广泛采用这种方法来减轻信息损失。然而，这种多扫描不可避免地增加了计算复杂性，特别是对于高分辨率输入。此外，第3节中的实证研究表明，这些多方向扫描之间也存在大量信息冗余。第三，我们在第3节中的发现表明，Mamba[21]在标记交互中容易出现长距离衰减，这意味着序列中距离较远的标记之间的交互减弱。因此，即使是先前扫描过的距离较远但相关的像素，查询像素也无法有效利用。
在这里插入图片描述

在这项工作中，我们提出MambaIRv2，旨在解决因果状态空间建模的不利影响。由于ViTs[17, 47]自然支持非因果处理，我们的关键思想是将类似ViT的非因果建模集成到基于Mamba的方法中。为此，我们首先深入研究注意力和状态空间之间的联系，以获取有价值的见解。我们在第4.1节中的深入分析表明，状态空间方程的输出矩阵类似于注意力机制中的查询。这种相似性启发我们利用输出矩阵来“查询”未扫描序列中的相关像素。受益于在扫描序列之外进行注意力聚焦，这种策略自然也消除了多方向扫描的需要。此外，为了促进远距离但相关像素之间的交互，我们提议重新构建图像，使相似像素在一维序列中在空间上更接近。通过这种方式，它允许进行语义而非空间序列建模，从而减轻长距离衰减的影响。由于所提出的方法使Mamba的行为类似于注意力机制，因此我们将其称为“注意力状态空间恢复”。

1.3 研究贡献

总体而言，我们做出了三个关键贡献：

提出注意力状态空间方程（ASE）：利用Mamba原始状态空间方程中的提示学习，在扫描序列之外查询语义相似的像素。具体来说，提示被设计为表示整个图像中相似的像素集，然后我们通过残差加法将代表性提示合并到状态空间方程的输出矩阵中，从而得到我们的ASE。作为核心组件，所提出的ASE不仅减轻了Mamba的因果性以提高性能，还实现了单遍扫描以提高效率。
开发语义引导邻域（SGN）机制：鼓励远距离但相似的像素之间进行强交互。具体而言，我们首先为每个像素分配相应的语义标签。然后，我们根据这些标签重新构建图像，以生成语义邻域一维序列，其中语义相似的像素在空间上也彼此靠近。由于减轻了Mamba的长距离衰减，SGN促进了原始图像中距离较远的像素之间的有效交互。
构建MambaIRv2模型：将两个核心模块和其他辅助部分集成在一起，我们提出了MambaIRv2，这是一种注意力状态空间恢复方法，为Mamba的状态空间建模赋予了类似ViT的非因果能力。大量实验表明，MambaIRv2显著提高了有效性和效率。特别是，MambaIRv2在Urban100数据集上，对于2倍轻量级超分辨率任务，比基于Transformer的最先进基线SRFormer[70]的PSNR高出0.35dB，在Manga109数据集上，对于2倍经典超分辨率任务，比HAT[8]的PSNR高出0.29dB。

2. 相关工作

2.1 基于CNN的图像恢复方法

近年来，图像恢复领域取得了巨大进展。早期尝试通常采用卷积神经网络（CNNs），例如用于图像超分辨率的SRCNN[15]、用于图像去噪的DnCNN[56]以及用于减少JPEG压缩伪影的ARCNN[16]。为了进一步提高基于CNN的方法的性能，人们引入了各种技术。例如，EDSR[29]采用残差连接策略，允许训练非常深的神经网络；RDN[65]使用密集连接来提高模型的表示能力；RCAN[64]引入通道注意力机制来选择显著通道，随后SAN[13]使用二阶注意力机制来提高性能。尽管基于CNN的方法取得了很大进展，但卷积算子本质上把感受野限制在局部内核，阻止了远距离像素之间的交互。

2.2 基于Transformer的图像恢复方法

由于Transformer[47]已在多个计算机视觉任务中证明了其有效性，因此将Transformer应用于图像恢复似乎很有前景。然而，直接应用普通自注意力机制，其计算复杂度与输入大小呈二次方关系，成本高昂且不切实际。为了提高注意力机制的效率，人们开发了多种技术。例如，IPT[6]将一幅图像划分为几个小补丁，并使用自注意力机制独立处理每个补丁。之后，SwinIR[33]进一步引入了移位窗口自注意力[36]来提高性能。ART[55]和OminiSR[48]利用稀疏注意力机制通过扩大注意力窗口来扩展感受野。GRL[32]采用锚点注意力机制来学习图像的局部、区域和全局层次结构。最近，ATD[61]使用自适应标记字典来存储与输入无关的知识，从而使注意力能够关注局部窗口之外的信息。

2.3 基于Mamba的图像恢复方法

为了平衡高效计算和全局感受野，最近人们在图像恢复中探索了Mamba[21]，并取得了有前景的结果。MambaIR[23]是最早将Mamba引入图像恢复的方法之一，它解决了两个特定挑战，即局部像素遗忘和通道冗余。从那时起，Mamba模型在各种图像恢复任务中得到了探索。FreqMamba[68]在傅里叶域中使用状态空间模型进行图像去雨，以感知全局退化；MambaLLIE[49]改进了状态空间方程，以实现低光图像增强任务中的局部性增强。此外，Mamba在图像去雾[69]、去模糊[20]和其他任务[3, 35, 42, 43, 50, 51, 62]中也取得了有前景的结果。然而，现有方法仍然受限于Mamba的因果建模本质。鉴于图像恢复是一个非因果任务，这种不匹配可能会导致性能受限和效率低下。
在这里插入图片描述

3. 研究动机

在这里插入图片描述

3.1 基于Mamba的图像恢复方法

现有的状态空间恢复方法主要是基于Mamba[21]架构开发的。形式上，Mamba采用离散状态空间方程来建模标记之间的交互：
$\begin{aligned} & h_{i}=\overline{A} h_{i - 1}+\overline{B} x_{i}, \\ & y_{i}=C h_{i}+D x_{i}, \end{aligned}$
其中 $\overline{A}=\exp (\Delta A)$ 是控制矩阵， $\overline{B}=(\Delta A)^{-1}(\exp (\Delta A)-I)\Delta B \approx \Delta B$ 是输入矩阵，C是输出矩阵。公式（1）表明，第 $i$ 个标记完全依赖于其前面的 $i - 1$ 个标记，即状态空间建模具有因果属性。尽管这种因果性对诸如自然语言处理等自回归任务有帮助，但它给图像恢复带来了挑战。

3.2 因果建模带来的挑战

现有的基于Mamba的方法通常采用特定的扫描策略将二维图像展开为一维序列，以便使用Mamba进行顺序建模。然而，第 $i$ 个像素只能看到整个图像中有限的 $i - 1$ 个像素，无法全局利用相似像素。为此，当前方法通常采用多方向扫描以获得更广泛的感受野，这不可避免地伴随着计算复杂度的增加。此外，如图2（a）所示，所有测试数据集上不同扫描序列的相似度甚至达到0.7以上，这表明存在高度相关性和大量冗余。而且，由于其因果性，Mamba[21]本身也存在长距离衰减缺陷。具体而言，像素之间的交互可以用控制矩阵 $\overline{A}^{k}$ 的幂来定量表示，其中 $k$ 是成对距离（证明见补充材料）。在图2（b）中，我们展示了学习到的 $\overline{A}$ 在统计上小于1。因此，当两个像素相距较远时，即 $k$ 较大时，交互 $\overline{A}^{k}$ 会变弱，这表明当前方法无法利用距离较远但有用的已扫描像素。

4. 注意力状态空间恢复

4.1 搭建注意力和状态空间的桥梁

如[24]所指出的，状态空间与注意力有很强的关系，这可能为将非因果建模能力融入Mamba提供见解。在本节中，我们首先将注意力和状态空间重新表述为通用形式以便比较，然后进行详细的联系分析。

4.1.1 注意力的重新表述

由于Mamba属于具有线性复杂度的因果模型，我们采用相应的因果线性注意力[28]作为其对应。具体而言，给定查询、键和值矩阵 $Q$ 、 $K$ 、 $v$ ，线性注意力的输出计算如下：
$y_{i}=\sum_{j = 1}^{i}\frac{Q_{i}K_{j}^{\top}}{\sum_{t = 1}^{i}Q_{i}K_{t}^{\top}}V_{j}=\frac{Q_{i}(\sum_{j = 1}^{i}K_{j}^{\top}V_{j})}{Q_{i}(\sum_{t = 1}^{i}K_{t}^{\top})}$