前段时间我们开发的基于Mamba的图像复原模型MambaIR受到了研究者的关注，并展示了Mamba作为除了CNN和Transformer之外的第三种图像修复backbone的潜力。

最近，我们在MambaIR的基础上继续开发了一个新版本MambaIRv2，目前已经被 CVPR 2025 接收。

下面是相关的论文和代码链接。

论文：arxiv.org/pdf/2411.15269

代码：github.com/csguoh/MambaIR

引言

简单来说，我们的MambaIRv2的主要亮点包括：

• 我们将原始MamabIR的扫描顺序从4次降低为了1次，同时不影响图像均匀采样
• 我们的MambaIRv2在lightSR任务上用更少参数实现了比SRformer更好的结果
• 在经典SR任务上，MambaIRv2用相同的参数量超过了HAT

Motivation

MambaIRv2要解决的核心问题就是Mamba模型的因果扫描局限性。具体来说，因果扫描就是如果我们将图像展开1D序列，那么在展开序列中的第i个像素只能利用到它的前i-1个像素，而无法用到序列之后的其他像素。这种mamba，或者更一般的自回归模型，的固有性质虽然对NLP任务很match，但是对非因果的视觉任务其实是存在局限的，因为大多数的视觉任务的所有token，即像素，是一次性可观的。

下面这个图总结了之前MambaIR的因果扫描 v.s. MambaIRv2的注意力扫描

因此，我们提出了MambaIRv2来实现类似ViT的那种第i个toekn可以一次性看到所有其他token的特性。也就是给Mamba加上类似注意力的功能，这也是为什么我们题目中有个”Attentive”的原因。

Some Findings

为了给Mamba加上注意力，我们首先观察了Mamba模型在图像复原任务中的特性：

1. 多方向信息冗余

目前的视觉Mamba方法，例如之前的MambaIR，主要是使用了多方向扫描来克服Mamba架构的因果扫描特性。虽然这样可以使得第i个像素看到序列中之后的像素，这也带来了高的计算复杂度。实验中，我们发现这种计算复杂度事实上是可以减少的，即不同方向的扫描序列之间其实是十分类似的，这也就为我们v2模型的单方向扫描奠定了基础。我们可视化了MambaIR的四个方向扫描得到的特征的相似度如下:

2. 远距离交互衰减

第二个问题则是远距离衰减，即第 i 个像素甚至不能很好地使用之前已经扫描过的但是距离较远的像素。具体来说，我们证明了Mamba的SSM建模中两个距离为k的像素之间的交互可以正比于 ，其中的 是原始状态空间方程的状态矩阵，见下面公式，而 则表示两个像素之间的距离。

在实际中我们发现，模型学到的 实际上显著地小于1，见下图，从而导致如果两个像素距离很远，即 很大时，两者之间的交互就会变得很弱。

3. 连接SSM和Attention

我们事实上可以在数学上将SSM类比为Attention。具体推导我就不写了(不是我懒)，这里直接放结论：

我们可以把线性Attention改写为以下形式：

同时，按照类似的形式，我们也可以把SSM等价地转化为以下形式：

通过对比上面两个不同形式，我们可以看出状态空间方程和注意力机制之间的联系如下：

因此，一个直观的思路就是重新修改SSM的矩阵C，来让它实现类似注意力机制中的全局查询功能，也就实现了将注意力机制加入到Mamba中。（上面这些特性启发了我们设计新的MamabIRv2模型，也可以作为大家的takeaway

方法

我们提出的MambaIRv2模型的主要架构如上所示。模型架构的核心组件是Attentive State Space Module(ASSM)。其中ASSM又进一步地包括了Attentive State-space Equation(ASE)和Semantic Guided Neighboring(SGN)。下面是具体的技术细节：

Attentive State Space Module

如图（a）所示，给定输入特征 ，其中 和 分别是高度和宽度， 是通道维度，我们首先对 应用位置编码以保留原始结构信息。之后，我们提出了Semantic Guided Neighboring（SGN）将2D图像展开成1D序列，以便后续的Attentive State－space Equation （ASE）建模。最后，另一个SGN作为之前操作的逆操作，将序列重新折叠回图像，然后通过线性投影得到模块输出。

Attentive State-space Equation

按照前面的Attention与SSM的关联性分享，我们的目标是修改输出矩阵 C以便全局查询图像中的相关像素。为此，我们提出了Attentive State－space Equation（ASE），它基于Mamba的原始状态空间方程，但具有非因果性质。如图（b）所示，我们提出的ASE将提示（prompts）纳入到C中，这些提示学习代表具有相似语义的一定像素集合，以补充未扫描像素缺失的信息。具体来说，我们首先构建提示池 来包含一系列可学习的典型prompts。对于 的参数化，我们采用低秩解耦以提高参数效率：

其中 在不同块之间共享， 是块特定的， r 是内部秩，且 。我们这里使用低秩解耦的主要思想是，我们希望不同的块共享类似的特征空间，即 是共享的，而共享特征的组合系数可以因不同的block而异，即 是特定的。之后，我们开发路由策略从 中选择以获得 个实例特定的提示 ，这些将被添加到 中以包含未扫描像素的信息。具体来说，给定展开的输入特征 ，我们使用线性层将 的通道维度从 投影到 ，然后通过 LogSoftmax预测对数概率，这表示 中每个提示被 采样的概率， 。之后，我们在对数概率上引入gumbel－softmax技巧，以允许可微分的提示选择操作，以获得路由矩阵 。然后，通过矩阵乘法生成实例特定的提示，即 。最后，我们通过残差加法将 纳入 中，形成所提出的注意力状态空间方程：

作为另一个优势，所提出的注意力状态空间方程允许模型只使用单一方向进行扫描，消除了现有方法中多方向扫描的高计算成本和冗余。

Semantic Guided Neighboring

Mamba的因果建模属性导致了长距离衰减的不利影响。在现有的基于Mamba的图像恢复方法中，原始图像中相距遥远的像素通常在展开的序列中仍然相距遥远，导致查询像素对于已经扫描的、空间上遥远但相似的像素的利用不足。为此，我们提出了Semantic Guided Neighboring（SGN），如图(c)所示。我们的关键见解是，与自回归语言建模不同，图像恢复是一个非因果任务，所有像素都可以一次性观察到，因此我们可以重新定义令牌邻域，使得语义相似的令牌在展开的序列中空间上更接近。按照这个想法，我们首先确定每个像素的语义标签。注意，在ASE中的路由矩阵，它已经学习了每个像素的提示类别，因此我们利用这种现成的语义来重构图像。具体来说，我们提出的SGN-unfold将具有第i个提示类别的像素组合在一起形成第i个语义组，然后根据类别值i组合不同的组以生成语义邻近序列。之后，我们将这个序列输入到提出的ASE中进行状态空间建模。最后，我们使用SGN-fold作为SGN-unfold的逆变换，将语义空间序列重新塑形回空间空间特征图以获得输出。

实验

轻量化SR

经典SR

JPEG伪影修复

图像去噪声

进一步讨论

与MambaIR-V1版本对比

得益于减少的扫描方向，与之前的在像素空间执行4次扫描的MambaIR相比，所提出的MambaIRv2的一个显著优势是它只需要在语义空间进行单次扫描。如下表所示，MambaIRv2不仅效率高，还提升了性能。例如，MambaIRv2在2×Urban100上减少了43% 的参数量和50% 的计算负担，同时仍然相比MambaIR取得了0.34dB 的PSNR增益。

Attentive Map Visualization

在提出的注意力状态空间方程中，提示prompts在表示整个图像中相似像素以促进查询像素看到扫描序列之外的部分中扮演了重要角色。如上图所示，可以看到查询像素通过提示得以关注图像中对应的语义部分，从而实现了类似注意力机制的全局信息聚合。

结论

在这项工作中，我们引入了MambaIRv2，通过解决Mamba架构的因果建模特性问题以增强现有的基于状态空间模型的图像复原算法。我们提出了Attentive State-space Equation（ASE）通过提示学习策略扩大了像素的感知范围，并且只需要进行单次扫描从而大大提升了运行效率。此外，我们还引入了Semantic Guided Neighboring（SGN），它通过将相似的像素放置得更近来处理长距离衰减问题。这些创新使得MambaIRv2能够将类似ViT的非因果能力整合到基于Mamba的模型中，实现注意力状态空间图像修复。广泛的实验证实了MambaIRv2是一个高效、高性能的图像恢复骨干网络。