輕量化MobileMamba視覺模型來了｜浙大/騰訊優(yōu)圖/華中科大聯(lián)合出品
浙大、騰訊優(yōu)圖、華中科技大學(xué)的團隊，提出輕量化MobileMamba！
既良好地平衡了效率與效果，推理速度遠超現(xiàn)有基于Mamba的模型。
一直以來，輕量化模型研究的主陣地都在CNN和Transformer的設(shè)計。
但CNN的局部有效感受野在高分辨率輸入時，難以獲得長距離依賴；盡管Transformer有著全局建模能力，但是其平方級計算復(fù)雜度，限制了其在高分辨率下的輕量化應(yīng)用。
最近的狀態(tài)空間模型如Mamba，因其線性計算復(fù)雜度和出色的效果被廣泛用在視覺領(lǐng)域。
然而，基于Mamba的輕量化模型雖然FLOPs低，但是實際的吞吐量極低。
△最近基于CNN/Transformer/Mamba方法的效果 vs. FLOPs對比
團隊首先在粗粒度上設(shè)計了三階段網(wǎng)絡(luò)顯著提升推理速度。
隨后在細粒度上提出了高效多感受野特征交互（MRFFI）模塊包含長距離小波變換增強Mamba （WTE-Mamba）、高效多核深度可分離卷積（MK-DeConv）和去冗余恒等映射三個部分——這有利于在長距離建模的特征上融合多尺度多感受野信息并加強高頻細節(jié)特征提取。
最后，使用兩個訓(xùn)練和一個推理策略，進一步提升模型的性能與效率。
大量實驗驗證，MobileMamba在ImageNet-1K數(shù)據(jù)集上的Top - 1準確率最高可達83.6，且速度是LocalVim的21倍、EfficientVMamba的3.3倍。
同時，大量的下游任務(wù)實驗也驗證了該方法在高分辨率輸入情況下，取得了效果與效率的最佳平衡。
現(xiàn)存缺陷：成本高、速度低
隨著移動設(shè)備的普及，資源受限環(huán)境中對高效、快速且準確的視覺處理需求日益增長。
開發(fā)輕量化模型，有助于顯著降低計算和存儲成本，還能提升推理速度，從而拓展技術(shù)的應(yīng)用范圍。
現(xiàn)有被廣泛研究的輕量化模型，主要被分為基于CNN和Transformer的結(jié)構(gòu)。
基于CNN的MobileNet，設(shè)計了深度可分離卷積大幅度減少了計算復(fù)雜度；GhostNet提出將原本將原本全通道1x1卷積替換為半數(shù)通道進行廉價計算，另半數(shù)通道直接恒等映射。
這些方法給后續(xù)基于CNN的工作奠定了良好的基礎(chǔ)。
但是基于CNN方法的主要缺陷在于其局部感受野，如圖(i)所示，其ERF僅在中間區(qū)域而缺少遠距離的相關(guān)性。
并且在下游任務(wù)高分辨率輸入下，基于CNN的方法僅能通過堆疊計算量來換取性能的少量提升。
如圖(ii)所示，ViT有著全局感受野和長距離建模能力。但由于其平方級別的計算復(fù)雜度，計算開銷比CNN更大。
一些工作嘗試從減少分辨率或者減少通道數(shù)上，來改減少所帶來的計算復(fù)雜度的增長，也取得了出色的效果。
不過，基于純ViT的結(jié)構(gòu)缺少了歸納偏置，因此，越來越多的研究者將CNN與Transformer結(jié)合得到混合結(jié)構(gòu)，獲得更好的效果，并獲得局部和全局的感受野（如圖(iii)）。
不過，尤其在下游任務(wù)高分辨率輸入下，基于ViT的方法仍然受到平方級別計算復(fù)雜度的問題。
提出MobileMamba
最近，由于狀態(tài)空間模型捕捉長距離依賴關(guān)系并且線性的計算復(fù)雜度表現(xiàn)出色，引起了廣泛關(guān)注，大量研究者將其應(yīng)用于視覺領(lǐng)域，效果和效率都取得了出色的效果。
基于Mamba的輕量化模型LocalMamba提出了將圖像劃分為窗口并在窗口內(nèi)局部掃描的方式減少計算復(fù)雜度，而EfficientVMamba設(shè)計了高效2D掃描方式來降低計算復(fù)雜度。
不過這兩種模型都僅公布了FLOPs，而FLOPs低并不能代表推理速度快。
經(jīng)實驗發(fā)現(xiàn)（圖2），現(xiàn)有的基于Mamba結(jié)構(gòu)的推理速度較慢并且效果較差。
MobileMamba團隊分別從粗粒度、細粒度和訓(xùn)練測試策略三個方面來設(shè)計高效輕量化網(wǎng)絡(luò)。
首先，研究人員討論了四階段和三階段在準確率、速度、FLOPs上的權(quán)衡。
在同等吞吐量下，三階段網(wǎng)絡(luò)會取得更高的準確率；同樣的相同效果下三階段網(wǎng)絡(luò)有著更高的吞吐量。
因此，團隊選擇三階段網(wǎng)絡(luò)作為MobileMamba的粗粒度設(shè)計框架。
在細粒度模塊設(shè)計方面，研究人員提出了高效高效多感受野特征交互 （MRFFI）模塊。
具體來說，將輸入特征根據(jù)通道維度劃分三個部分。
第一部分將通過小波變換增強的Mamba模塊提取全局特征的同時加強邊緣細節(jié)等細粒度信息的提取能力。
第二部分通過高效多核深度可分離卷積操作獲取多尺度感受野的感知能力。
然后部分通過去冗余恒等映射，減少高維空間下通道冗余的問題，并減少計算復(fù)雜度提高運算速度。
最終經(jīng)過MRFFI得到的特征融合了全局和多尺度局部的多感受野信息，并且加強了邊緣細節(jié)的高頻信息提取能力。
最后，研究人員通過兩個訓(xùn)練階段策略知識蒸餾和延長訓(xùn)練輪數(shù)增強模型的學(xué)習(xí)能力，提升模型效果；以及一個歸一化層融合的測試階段策略提升模型的推理速度。
△MobileMamba結(jié)構(gòu)概述
實驗結(jié)果
實驗表明，MobileMamba有著全局感受野的同時，高效多核深度可分離卷積操作有助于提取相鄰信息。
通過與SoTA方法的對比可知，MobileMamba從200M到4G FLOPs的模型在使用訓(xùn)練策略后，在ImageNet-1K上的Top-1，分別達到76.9、78.9、80.7、82.2、83.3、83.6效果，均超過現(xiàn)有基于CNN、ViT和Mamba的方法。
與同為Mamba的方法相比，MobileMamba比LocalVim在Top-1上提升0.7↑的同時，速度快21倍；比EfficientVMamba提升2.0↑的同時速度快3.3↑倍。
這均顯著優(yōu)于現(xiàn)有基于Mamba的輕量化模型設(shè)計。
同時，在下游任務(wù)目標(biāo)檢測、實力分割、語義分割上大量實驗上也驗證了方法的有效性。
在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。
在高分辨率輸入的下游任務(wù)與基于CNN的MobileNetv2和ViT的MobileViTv2相比分別提升7.2↑和0.4↑，并且FLOPs僅有其8.5%和11.2%。
總的來說，MobileMamba貢獻如下：
- 提出了一個輕量級的三階段MobileMamba框架，該框架在性能和效率之間實現(xiàn)了良好的平衡。MobileMamba的有效性和效率已經(jīng)在分類任務(wù)以及三個高分辨率輸入的下游任務(wù)中得到了驗證。
- 設(shè)計了一個高效的多感受野特征交互（MRFFI）模塊，以通過更大的有效感受野增強多尺度感知能力，并改進細粒度高頻邊緣信息的提取。
- MobileMamba通過在不同F(xiàn)LOPs大小的模型上采用訓(xùn)練和測試策略，顯著提升了性能和效率。
