CVPR2025|清華提出MambaIRv2：影像復原超強SOTA模型

2025-04-21 11:21 CVer

點選下方卡片，關注“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公眾號
整理不易，請點贊和在看


相關文章