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

最新推荐文章于 2025-04-21 16:31:59 发布

Amusi（CVer）

最新推荐文章于 2025-04-21 16:31:59 发布

阅读量500

点赞数

文章标签：人工智能

原文链接：https://mp.weixin.qq.com/s?__biz=MzUxNjcxMjQxNg==&mid=2247617323&idx=2&sn=cffccae87d5b6705f1d41d1ac75f47ae&chksm=f8d95425746c448d8ee8d4d2c4d5b526aac95d06e252a1d4134bfc3cd4bfe4350e73907d7135&scene=126&sessionid=0

版权

点击下方卡片，关注“CVer”公众号

AI/CV重磅干货，第一时间送达

点击进入—>【顶会/顶刊】投稿交流群

添加微信号：CVer2233，小助手会拉你进群！

扫描下方二维码，加入CVer学术星球！可以获得最新顶会/顶刊上的论文idea和CV从入门到精通资料，及最前沿应用！发论文/搞科研/涨薪，强烈推荐！

作者：郭航

https://zhuanlan.zhihu.com/p/9375585949

前段时间我们开发的基于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

如图(a)所示，给定输入特征 x∈RH×W×C ，其中 H 和 W 分别是高度和宽度， C 是通道维度，我们首先对 x 应用位置编码以保留原始结构信息。之后，我们提出了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中，这些提示学习代表具有相似语义的一定像素集合，以补充未扫描像素缺失的信息。具体来说，我们首先构建提示池 P来包含一系列可学习的典型prompts。对于 P 的参数化，我们采用低秩解耦以提高参数效率：

Semantic Guided Neighboring

Mamba的因果建模属性导致了长距离衰减的不利影响。在现有的基于Mamba的图像恢复方法中，原始图像中相距遥远的像素通常在展开的序列中仍然相距遥远，导致查询像素对于已经扫描的、空间上遥远但相似的像素的利用不足。为此，我们提出了Semantic Guided Neighboring（SGN），如图(c)所示。我们的关键见解是，与自回归语言建模不同，图像恢复是一个非因果任务，所有像素都可以一次性观察到，因此我们可以重新定义令牌邻域，使得语义相似的令牌在展开的序列中空间上更接近。按照这个想法，我们首先确定每个像素的语义标签。注意，在ASE中的路由矩阵 R ，它已经学习了每个像素的提示类别，因此我们利用这种现成的语义来重构图像。具体来说，我们提出的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是一个高效、高性能的图像恢复骨干网络。

何恺明在MIT授课的课件PPT下载

在CVer公众号后台回复：何恺明，即可下载本课程的所有566页课件PPT！赶紧学起来！

CVPR 2025 论文和代码下载

在CVer公众号后台回复：CVPR2025，即可下载CVPR 2025论文和代码开源的论文合集

ECCV 2024 论文和代码下载

在CVer公众号后台回复：ECCV2024，即可下载ECCV 2024论文和代码开源的论文合集

CV垂直方向和论文投稿交流群成立

扫描下方二维码，或者添加微信号：CVer2233，即可添加CVer小助手微信，便可申请加入CVer-垂直方向和论文投稿微信交流群。另外其他垂直方向已涵盖：目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch、TensorFlow和Transformer、NeRF、3DGS、Mamba等。
一定要备注：研究方向+地点+学校/公司+昵称（如Mamba、多模态学习或者论文投稿+上海+上交+卡卡），根据格式备注，可更快被通过且邀请进群

▲扫码或加微信号: CVer2233，进交流群
CVer计算机视觉（知识星球）来了！想要了解最新最快最好的CV/DL/AI论文速递、优质实战项目、AI行业前沿、从入门到精通学习教程等资料，欢迎扫描下方二维码，加入CVer计算机视觉（知识星球），已汇集上万人！

▲扫码加入星球学习

▲点击上方卡片，关注CVer公众号
整理不易，请点赞和在看