論文地址：https://arxiv.org/pdf/2506.15976
代碼地址：https://github.com/CiaoHe/bi-mamba

---

關注UP CV縫合怪，分享最計算機視覺新即插即用模塊，并提供配套的論文資料與代碼。
https://space.bilibili.com/473764881

摘要

Mamba 是一種狀態空間模型 (SSM)，它通過將遞歸重鑄為并行選擇性掃描來加速訓練，最近已成為一種線性縮放、高效的自注意力替代方案。由于其單向性，Mamba 中的每個狀態都只包含其先前狀態的信息，而對之后的狀態視而不見。當前，基于 Mamba 的計算機視覺方法通常通過在 Mamba 的全局前向掃描基礎上增加全局后向掃描來克服這一限制，形成雙向掃描以恢復全部感受野。然而，此操作會使計算負荷加倍，從而削弱 Mamba 原本的效率優勢。

為了消除這些額外的掃描，本研究引入了局部雙向 SSM 模塊 LBMamba，該模塊在前向選擇性掃描內嵌入輕量級局部后向掃描，并在每線程寄存器中完全執行。在本研究在 LBMamba 的基礎上提出了可擴展的視覺骨干網絡 LBVim，它每兩層交替掃描方向，無需額外的后向掃描即可恢復全局感受野。本研究在自然圖像和全幻燈片圖像 (WSI) 上驗證了該方法的多功能性。

結果表明，LBVim 始終提供優越的性能-吞吐量權衡。也就是說，在相同的吞吐量下，LBVim 在 ImageNet-1K 分類數據集上的 top-1 準確率提高了 0.8% 到 1.6%，在 ADE20K 語義分割數據集上的 mIoU 提高了 0.6% 到 2.7%，在 COCO 檢測數據集上的 APb 和 APm 分別提高了 0.9% 和 1.1%。本研究還將 LBMamba 集成到最先進的病理學多實例學習 (MIL) 方法 MambaMIL 中，該方法使用單向掃描。在 3 個公開的 WSI 分類數據集上進行的實驗表明，本研究的方法實現了高達 3.06% 的 AUC、3.39% 的 F1 和 1.67% 的準確率的相對改進。

引言

本研究關注基于Mamba的狀態空間模型（SSM）在計算機視覺中的應用。SSM因其線性時間和內存復雜度而成為序列建模中自注意力機制的替代方案。然而，傳統SSM的訓練和推理速度仍然受到限制，因為它們無法充分利用現代GPU的并行性。Mamba通過將狀態更新與隱藏到輸出的卷積解耦，并將計算重新表述為在現代GPU上高效運行的并行選擇性掃描來克服這個問題。因此，Mamba在長距離任務上達到了與Transformer相當的精度，同時展現出更好的分辨率縮放特性，使其成為研究和生產系統的理想選擇。它最初是為自然語言處理引入的，后來被應用于計算機視覺。基于Mamba的選擇性掃描內核構建的視覺模型可顯著提高GPU速度和節省內存，同時始終優于基于Transformer的基線模型。

然而，標準的基于Mamba的計算機視覺模型通常需要從不同方向多次掃描圖像才能提高性能。這種多次掃描有兩個主要原因：第一，Mamba將圖像視為扁平化的一維序列，因此單一的從左到右的掃描只能捕獲行方向的上下文信息。為了恢復垂直方向的依賴關系，視覺流水線通常會對圖像塊的列方向進行額外的掃描，從而產生兩次正交掃描，共同逼近二維空間關系。盡管一些專門的二維Mamba/SSM方法被提出以更直接地解決這種結構不匹配問題，但多次掃描的需求依然存在。第二個原因是SSM的單向性: 位置t的潛在狀態僅以先前位置1到t為條件。因此，模型無法獲取位置t之后的信息，這通常會導致視覺任務的性能欠佳。

目前常見的解決方案是添加反向掃描（從右到左或從下到上）以恢復對未來標記的訪問，從而產生雙向掃描機制。雖然這種策略重建了完整的感受野，但每次額外的掃描都會大致使計算量加倍，從而削弱了Mamba最初提供的效率優勢。為了解決全局雙向掃描帶來的額外計算開銷問題，本研究提出了LBMamba，一種局部雙向SSM，以及用于視覺任務的LBVim框架。LBMamba在正向選擇性掃描內部嵌入了輕量級的局部反向掃描，并在每個線程的寄存器中完全執行，從而無需進行全局反向掃描。基于LBMamba構建的LBVim視覺主干網絡每兩層交替掃描方向，無需額外的反向掃描即可恢復全局感受野。

論文創新點

本研究提出了一個名為LBMamba的局部雙向SSM模塊，旨在提高Mamba模型在計算機視覺任務中的效率和性能。其創新點如下：

1. ? 引入了局部反向掃描機制： ?

   • 不同于傳統的全局雙向掃描，LBMamba在Mamba的前向選擇性掃描過程中嵌入了輕量級的局部反向掃描。
   • 這種局部反向掃描在子序列內進行，無需進行全局反向掃描，從而顯著提高了計算效率。

2. ?? 基于硬件感知的線程級雙向掃描算子： ??

   • 本研究設計了一個快速的、硬件感知的線程級雙向掃描算子。
   • 該算子完全在線程私有寄存器中執行局部反向掃描，避免了額外的內存讀取和線程間通信，最大限度地減少了計算開銷。

3. 🚀 消除了全局反向掃描的額外開銷： 🚀

   • 通過將局部反向掃描整合到前向掃描過程中，LBMamba消除了全局反向掃描帶來的額外計算負擔和內存訪問，從而顯著提高了模型的吞吐量。
   • 實驗結果表明，LBMamba的運行時間僅增加了2%，卻避免了全局反向掃描帶來的運行時間翻倍。

4. 🔄 基于序列反轉操作恢復全局感受野： 🔄

   • 為了彌補局部反向掃描在全局感受野方面的不足，本研究提出在每個編碼器層后反轉特征序列。
   • 這種交替的掃描方向策略，使得每個token在每兩個編碼器層后都能獲得全局感受野，從而增強了模型對長距離上下文信息的建模能力。

5. ? 在不同視覺任務和數據集上的有效性驗證： ?

   • 本研究在自然圖像和病理圖像等不同類型的圖像數據上進行了廣泛的實驗，涵蓋了圖像分類、語義分割、目標檢測和多實例學習等多種計算機視覺任務。
   • 實驗結果表明，LBMamba在多種任務和數據集上均取得了顯著的性能提升，證明了其有效性和泛化能力。
   • 例如，在ImageNet-1K分類數據集上，相同吞吐量下，LBVim的top-1準確率比基線模型提高了0.8%到1.6%。在ADE20K語義分割數據集上，mIoU提高了0.6%到2.7%。

論文實驗