浙大、騰訊優(yōu)圖、華中科技大學(xué)的團(tuán)隊(duì)，提出輕量化MobileMamba！

既良好地平衡了效率與效果，推理速度遠(yuǎn)超現(xiàn)有基于Mamba的模型。

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

但CNN的局部有效感受野在高分辨率輸入時(shí)，難以獲得長距離依賴；盡管Transformer有著全局建模能力，但是其平方級(jí)計(jì)算復(fù)雜度，限制了其在高分辨率下的輕量化應(yīng)用。

最近的狀態(tài)空間模型如Mamba，因其線性計(jì)算復(fù)雜度和出色的效果被廣泛用在視覺領(lǐng)域。

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

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

團(tuán)隊(duì)首先在粗粒度上設(shè)計(jì)了三階段網(wǎng)絡(luò)顯著提升推理速度。

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

最后，使用兩個(gè)訓(xùn)練和一個(gè)推理策略，進(jìn)一步提升模型的性能與效率。

大量實(shí)驗(yàn)驗(yàn)證，MobileMamba在ImageNet-1K數(shù)據(jù)集上的Top - 1準(zhǔn)確率最高可達(dá)83.6，且速度是LocalVim的21倍、EfficientVMamba的3.3倍。

同時(shí)，大量的下游任務(wù)實(shí)驗(yàn)也驗(yàn)證了該方法在高分辨率輸入情況下，取得了效果與效率的最佳平衡。

現(xiàn)存缺陷：成本高、速度低

隨著移動(dòng)設(shè)備的普及，資源受限環(huán)境中對(duì)高效、快速且準(zhǔn)確的視覺處理需求日益增長。

開發(fā)輕量化模型，有助于顯著降低計(jì)算和存儲(chǔ)成本，還能提升推理速度，從而拓展技術(shù)的應(yīng)用范圍。

現(xiàn)有被廣泛研究的輕量化模型，主要被分為基于CNN和Transformer的結(jié)構(gòu)。

基于CNN的MobileNet，設(shè)計(jì)了深度可分離卷積大幅度減少了計(jì)算復(fù)雜度；GhostNet提出將原本將原本全通道1x1卷積替換為半數(shù)通道進(jìn)行廉價(jià)計(jì)算，另半數(shù)通道直接恒等映射。

這些方法給后續(xù)基于CNN的工作奠定了良好的基礎(chǔ)。

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

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

如圖(ii)所示，ViT有著全局感受野和長距離建模能力。但由于其平方級(jí)別的計(jì)算復(fù)雜度，計(jì)算開銷比CNN更大。

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

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

不過，尤其在下游任務(wù)高分辨率輸入下，基于ViT的方法仍然受到平方級(jí)別計(jì)算復(fù)雜度的問題。

提出MobileMamba

最近，由于狀態(tài)空間模型捕捉長距離依賴關(guān)系并且線性的計(jì)算復(fù)雜度表現(xiàn)出色，引起了廣泛關(guān)注，大量研究者將其應(yīng)用于視覺領(lǐng)域，效果和效率都取得了出色的效果。

基于Mamba的輕量化模型LocalMamba提出了將圖像劃分為窗口并在窗口內(nèi)局部掃描的方式減少計(jì)算復(fù)雜度，而EfficientVMamba設(shè)計(jì)了高效2D掃描方式來降低計(jì)算復(fù)雜度。

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

經(jīng)實(shí)驗(yàn)發(fā)現(xiàn)（圖2），現(xiàn)有的基于Mamba結(jié)構(gòu)的推理速度較慢并且效果較差。

MobileMamba團(tuán)隊(duì)分別從粗粒度、細(xì)粒度和訓(xùn)練測(cè)試策略三個(gè)方面來設(shè)計(jì)高效輕量化網(wǎng)絡(luò)。

首先，研究人員討論了四階段和三階段在準(zhǔn)確率、速度、FLOPs上的權(quán)衡。

在同等吞吐量下，三階段網(wǎng)絡(luò)會(huì)取得更高的準(zhǔn)確率；同樣的相同效果下三階段網(wǎng)絡(luò)有著更高的吞吐量。

因此，團(tuán)隊(duì)選擇三階段網(wǎng)絡(luò)作為MobileMamba的粗粒度設(shè)計(jì)框架。

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

具體來說，將輸入特征根據(jù)通道維度劃分三個(gè)部分。

第一部分將通過小波變換增強(qiáng)的Mamba模塊提取全局特征的同時(shí)加強(qiáng)邊緣細(xì)節(jié)等細(xì)粒度信息的提取能力。

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

然后部分通過去冗余恒等映射，減少高維空間下通道冗余的問題，并減少計(jì)算復(fù)雜度提高運(yùn)算速度。

最終經(jīng)過MRFFI得到的特征融合了全局和多尺度局部的多感受野信息，并且加強(qiáng)了邊緣細(xì)節(jié)的高頻信息提取能力。

最后，研究人員通過兩個(gè)訓(xùn)練階段策略知識(shí)蒸餾和延長訓(xùn)練輪數(shù)增強(qiáng)模型的學(xué)習(xí)能力，提升模型效果；以及一個(gè)歸一化層融合的測(cè)試階段策略提升模型的推理速度。

△MobileMamba結(jié)構(gòu)概述

實(shí)驗(yàn)結(jié)果

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

通過與SoTA方法的對(duì)比可知，MobileMamba從200M到4G FLOPs的模型在使用訓(xùn)練策略后，在ImageNet-1K上的Top-1，分別達(dá)到76.9、78.9、80.7、82.2、83.3、83.6效果，均超過現(xiàn)有基于CNN、ViT和Mamba的方法。

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

這均顯著優(yōu)于現(xiàn)有基于Mamba的輕量化模型設(shè)計(jì)。

同時(shí)，在下游任務(wù)目標(biāo)檢測(cè)、實(shí)力分割、語義分割上大量實(shí)驗(yàn)上也驗(yàn)證了方法的有效性。

在Mask RCNN上比EMO提升1.3↑在mAP并且吞吐量提升56%↑。

在RetinaNet上比EfficientVMamba提升+2.1↑在mAP并且吞吐量提升4.3↑倍。

在SSDLite通過提高分辨率達(dá)到24.0/29.5的mAP。

在DeepLabv3，Semantic FPN，and PSPNet上有著較少的FLOPs分別最高達(dá)到37.4/42.7/36.9的mIoU。

在高分辨率輸入的下游任務(wù)與基于CNN的MobileNetv2和ViT的MobileViTv2相比分別提升7.2↑和0.4↑，并且FLOPs僅有其8.5%和11.2%。

總的來說，MobileMamba貢獻(xiàn)如下：

• 提出了一個(gè)輕量級(jí)的三階段MobileMamba框架，該框架在性能和效率之間實(shí)現(xiàn)了良好的平衡。MobileMamba的有效性和效率已經(jīng)在分類任務(wù)以及三個(gè)高分辨率輸入的下游任務(wù)中得到了驗(yàn)證。
• 設(shè)計(jì)了一個(gè)高效的多感受野特征交互（MRFFI）模塊，以通過更大的有效感受野增強(qiáng)多尺度感知能力，并改進(jìn)細(xì)粒度高頻邊緣信息的提取。
• MobileMamba通過在不同F(xiàn)LOPs大小的模型上采用訓(xùn)練和測(cè)試策略，顯著提升了性能和效率。

論文鏈接：https://arxiv.org/pdf/2411.15941

項(xiàng)目代碼：https://github.com/lewandofskee/MobileMamba