輕量化MobileMamba視覺模型來了

浙大、騰訊優圖、華中科技大學的團隊，提出輕量化 MobileMamba！

既良好地平衡了效率與效果，推理速度遠超現有基于 Mamba 的模型。

一直以來，輕量化模型研究的主陣地都在 CNN 和 Transformer 的設計。

但 CNN 的局部有效感受野在高分辨率輸入時，難以獲得長距離依賴；盡管 Transformer 有着全局建模能力，但是其平方級計算複雜度，限制了其在高分辨率下的輕量化應用。

最近的狀态空間模型如 Mamba，因其線性計算複雜度和出色的效果被廣泛用在視覺領域。

然而，基于 Mamba 的輕量化模型雖然 FLOPs 低，但是實際的吞吐量極低。

△最近基于 CNN/Transformer/Mamba 方法的效果 vs. FLOPs 對比

團隊首先在粗粒度上設計了三階段網絡顯著提升推理速度。

随後在細粒度上提出了高效多感受野特征交互（MRFFI）模塊包含長距離小波變換增強 Mamba （WTE-Mamba）、高效多核深度可分離卷積（MK-DeConv）和去冗餘恒等映射三個部分——這有利于在長距離建模的特征上融合多尺度多感受野信息并加強高頻細節特征提取。

最後，使用兩個訓練和一個推理策略，進一步提升模型的性能與效率。

大量實驗驗證，MobileMamba 在 ImageNet-1K 數據集上的 Top - 1 準确率最高可達 83.6，且速度是 LocalVim 的 21 倍、EfficientVMamba 的 3.3 倍。

同時，大量的下遊任務實驗也驗證了該方法在高分辨率輸入情況下，取得了效果與效率的最佳平衡。

現存缺陷：成本高、速度低

随着移動設備的普及，資源受限環境中對高效、快速且準确的視覺處理需求日益增長。

開發輕量化模型，有助于顯著降低計算和存儲成本，還能提升推理速度，從而拓展技術的應用範圍。

現有被廣泛研究的輕量化模型，主要被分爲基于 CNN 和 Transformer 的結構。

基于 CNN 的 MobileNet，設計了深度可分離卷積大幅度減少了計算複雜度；GhostNet 提出将原本将原本全通道 1x1 卷積替換爲半數通道進行廉價計算，另半數通道直接恒等映射。

這些方法給後續基于 CNN 的工作奠定了良好的基礎。

但是基于 CNN 方法的主要缺陷在于其局部感受野，如圖 ( i ) 所示，其 ERF 僅在中間區域而缺少遠距離的相關性。

并且在下遊任務高分辨率輸入下，基于 CNN 的方法僅能通過堆疊計算量來換取性能的少量提升。

如圖 ( ii ) 所示，ViT 有着全局感受野和長距離建模能力。但由于其平方級别的計算複雜度，計算開銷比 CNN 更大。

一些工作嘗試從減少分辨率或者減少通道數上，來改減少所帶來的計算複雜度的增長，也取得了出色的效果。

不過，基于純 ViT 的結構缺少了歸納偏置，因此，越來越多的研究者将 CNN 與 Transformer 結合得到混合結構，獲得更好的效果，并獲得局部和全局的感受野（如圖 ( iii ) ）。

不過，尤其在下遊任務高分辨率輸入下，基于 ViT 的方法仍然受到平方級别計算複雜度的問題。

提出 MobileMamba

最近，由于狀态空間模型捕捉長距離依賴關系并且線性的計算複雜度表現出色，引起了廣泛關注，大量研究者将其應用于視覺領域，效果和效率都取得了出色的效果。

基于 Mamba 的輕量化模型 LocalMamba 提出了将圖像劃分爲窗口并在窗口内局部掃描的方式減少計算複雜度，而 EfficientVMamba 設計了高效 2D 掃描方式來降低計算複雜度。

不過這兩種模型都僅公布了 FLOPs，而FLOPs 低并不能代表推理速度快。

經實驗發現（圖 2），現有的基于 Mamba 結構的推理速度較慢并且效果較差。

MobileMamba 團隊分别從粗粒度、細粒度和訓練測試策略三個方面來設計高效輕量化網絡。

首先，研究人員讨論了四階段和三階段在準确率、速度、FLOPs 上的權衡。

在同等吞吐量下，三階段網絡會取得更高的準确率；同樣的相同效果下三階段網絡有着更高的吞吐量。

因此，團隊選擇三階段網絡作爲 MobileMamba 的粗粒度設計框架。

在細粒度模塊設計方面，研究人員提出了高效高效多感受野特征交互（MRFFI）模塊。

具體來說，将輸入特征根據通道維度劃分三個部分。

第一部分将通過小波變換增強的 Mamba 模塊提取全局特征的同時加強邊緣細節等細粒度信息的提取能力。

第二部分通過高效多核深度可分離卷積操作獲取多尺度感受野的感知能力。

然後部分通過去冗餘恒等映射，減少高維空間下通道冗餘的問題，并減少計算複雜度提高運算速度。

最終經過 MRFFI 得到的特征融合了全局和多尺度局部的多感受野信息，并且加強了邊緣細節的高頻信息提取能力。

最後，研究人員通過兩個訓練階段策略知識蒸餾和延長訓練輪數增強模型的學習能力，提升模型效果；以及一個歸一化層融合的測試階段策略提升模型的推理速度。

△MobileMamba 結構概述實驗結果

實驗表明，MobileMamba 有着全局感受野的同時，高效多核深度可分離卷積操作有助于提取相鄰信息。

通過與 SoTA 方法的對比可知，MobileMamba 從 200M 到 4G FLOPs 的模型在使用訓練策略後，在 ImageNet-1K 上的 Top-1，分别達到 76.9、78.9、80.7、82.2、83.3、83.6 效果，均超過現有基于 CNN、ViT 和 Mamba 的方法。

與同爲 Mamba 的方法相比，MobileMamba 比 LocalVim 在 Top-1 上提升 0.7 ↑的同時，速度快 21 倍；比 EfficientVMamba 提升 2.0 ↑的同時速度快 3.3 ↑倍。

這均顯著優于現有基于 Mamba 的輕量化模型設計。