輕量化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的輕量化模型設計。
同時，在下游任務目標檢測、實力分割、語義分割上大量實驗上也驗證了方法的有效性。
在Mask RCNN上比EMO提升1.3↑在mAP并且吞吐量提升56%↑。
在RetinaNet上比EfficientVMamba提升+2.1↑在mAP并且吞吐量提升4.3↑倍。
在SSDLite通過提高分辨率達到24.0/29.5的mAP。
在DeepLabv3，Semantic FPN，and PSPNet上有著較少的FLOPs分別最高達到37.4/42.7/36.9的mIoU。
在高分辨率輸入的下游任務與基于CNN的MobileNetv2和ViT的MobileViTv2相比分別提升7.2↑和0.4↑，并且FLOPs僅有其8.5%和11.2%。
總的來說，MobileMamba貢獻如下：
- 提出了一個輕量級的三階段MobileMamba框架，該框架在性能和效率之間實現了良好的平衡。MobileMamba的有效性和效率已經在分類任務以及三個高分辨率輸入的下游任務中得到了驗證。
- 設計了一個高效的多感受野特征交互（MRFFI）模塊，以通過更大的有效感受野增強多尺度感知能力，并改進細粒度高頻邊緣信息的提取。
- MobileMamba通過在不同FLOPs大小的模型上采用訓練和測試策略，顯著提升了性能和效率。
