CVPR 2025 | 清华提出MambaIRv2：图像复原超强SOTA模型

数据派THU

于 2025-04-17 17:02:10 发布

阅读量242

点赞数

原文链接：https://mp.weixin.qq.com/s?__biz=MzI1MjQ2OTQ3Ng==&mid=2247655680&idx=2&sn=f71d0db3fd982be58549eb1ea3fd5e13&chksm=e82fd375a3d4acb9bc5d180e414d92afaa7e0c0741578d54445eda35a22adfc111e75f341a4c&scene=126&sessionid=0

版权

来源：人工智能前沿讲习‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍
本文共1700字，建议阅读7分钟广泛的实验证实了MambaIRv2是一个高效、高性能的图像恢复骨干网络。

前段时间我们开发的基于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. 远距离交互衰减

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

Attentive State-space Equation

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

Semantic Guided Neighboring

实验

轻量化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是一个高效、高性能的图像恢复骨干网络。

编辑：黄继彦‍‍‍

关于我们

数据派THU作为数据科学类公众号，背靠清华大学大数据研究中心，分享前沿数据科学与大数据技术创新研究动态、持续传播数据科学知识，努力建设数据人才聚集平台、打造中国大数据最强集团军。