ViG基于視覺Mamba的通用視覺主干網(wǎng)絡(luò)

Vision Mamba的成功預(yù)示著將視覺表征學(xué)習(xí)轉(zhuǎn)換為線性復(fù)雜度視覺序列表征學(xué)習(xí)具有巨大的潛力。盡管以Vision Mamba為代表的線性視覺序列表征學(xué)習(xí)結(jié)構(gòu)在高清圖像上展示了顯著的效率優(yōu)勢，但在更為常見的分辨率，仍略遜于Transformer和CNN。

為了進(jìn)一步推進(jìn)線性復(fù)雜度視覺序列表征學(xué)習(xí)的效率，我們結(jié)合當(dāng)今時代計算設(shè)備的硬件特性去設(shè)計結(jié)構(gòu)，將自然語言序列建模中硬件計算更為友好的門控線性注意力模塊GLA引入到視覺表征學(xué)習(xí)中，進(jìn)一步設(shè)計參數(shù)高效的雙向建模、長短上下文的動態(tài)門控機(jī)制和硬件感知的雙向算子設(shè)計。通過以上設(shè)計，我們提出的ViG模型在各種任務(wù)上，精度、參數(shù)量、效率均優(yōu)于主流的Transformer和CNN模型。

? 論文鏈接：

https://arxiv.org/abs/2405.18425

? 本工作已被AAAI 2025收錄。

Vision Mamba的成功預(yù)示著將視覺表征學(xué)習(xí)轉(zhuǎn)換為線性復(fù)雜度視覺序列表征學(xué)習(xí)具有巨大的潛力。盡管以Vision Mamba為代表的線性視覺序列表征學(xué)習(xí)結(jié)構(gòu)在高清圖像上展示了顯著的效率優(yōu)勢，但在更為常見的分辨率，仍略遜于Transformer和CNN。

為了進(jìn)一步推進(jìn)線性復(fù)雜度視覺序列表征學(xué)習(xí)的效率，我們結(jié)合當(dāng)今時代計算設(shè)備的硬件特性去設(shè)計結(jié)構(gòu)，將自然語言序列建模中對硬件計算更為友好的門控線性注意力模塊(Gated Linear Attention, GLA)引入到視覺表征學(xué)習(xí)中，并且進(jìn)一步設(shè)計參數(shù)高效的雙向建模，長短上下文的動態(tài)門控機(jī)制和硬件感知的雙向算子設(shè)計。通過以上設(shè)計，我們提出的ViG模型在各種任務(wù)上，精度、參數(shù)量、效率均優(yōu)于主流的Transformer和CNN模型。

背景介紹

Transformer的自注意力（Self-Attention, SA）機(jī)制需要對歷史所有的輸入進(jìn)行交互，導(dǎo)致計算量隨著輸入序列長度呈現(xiàn)二次方增長：

而線性注意力(Linear Attention, LA)機(jī)制則將注意力機(jī)制中的softmax算子替換為簡單的矩陣相乘，從而能夠通過矩陣乘法的結(jié)合律，先進(jìn)行KV的計算從而將計算量壓縮為線性：

? 線性門控注意力(Gated Linear Attention, GLA)機(jī)制則進(jìn)一步引入門控機(jī)制控制對歷史信息的遺忘和更新從而提高了原有線性注意力機(jī)制的表征能力： ?

方法概述

雙向門控注意力機(jī)制

GLA原是用于處理一維時序信號如文本，圖片作為二維信號，在一維序列表達(dá)上具有多向的特點，針對這一特點，我們在原GLA的設(shè)計中，通過僅僅引入雙向門控設(shè)計，便能夠極大地提升對視覺信號的空間表征能力。 ? 我們提出的BiGLA算子將前向和反向的視覺序列壓縮到固定大小的隱狀態(tài)和中去。 ?

硬件感知的雙向算子設(shè)計

為了進(jìn)一步提升效率，我們設(shè)計了一個硬件感知的雙向?qū)崿F(xiàn)，將BiGLA的前向掃描和后向掃描合并到一個Triton算子中，無需實例化反向序列，只需要維護(hù)一個單向視覺序列，便能夠進(jìn)行多向掃描和融合，從而大幅減小顯存占用，并提升硬件運(yùn)行速度。

長短上下文的動態(tài)門控制機(jī)

BiGLA算子中雙向隱狀態(tài)和是沿著一維視覺序列進(jìn)行長上下文全局壓縮，為了進(jìn)一步加強(qiáng)對圖片二維空間細(xì)節(jié)的感知，我們引入了短上下文的卷積門控設(shè)計： ?

結(jié)構(gòu)設(shè)計

圍繞上述的基礎(chǔ)ViG模塊我們進(jìn)一步搭建兩個結(jié)構(gòu)變種：類似于Vision Transformer的樸素直筒結(jié)構(gòu)ViG和類似于CNN的層次化金字塔結(jié)構(gòu)ViG-H。

實驗結(jié)果

上圖顯示我們提出的ViG和ViG-H在參數(shù)量和精度的權(quán)衡上打敗了先進(jìn)的基于Transformer和CNN的模型。特別值得注意的是，我們基礎(chǔ)ViG模塊同時具備全局感受野和線性復(fù)雜度，這是之前CNN，基于原始注意力機(jī)制Transformer和基于窗口化注意力機(jī)制Transformer所達(dá)不到的。

上圖兩個表格也是顯示我們在速度上能夠和當(dāng)前先進(jìn)模型在224x224的圖片輸入上匹配。通過綜合對比我們可以看到，ViG-S 與 DeiT-B精度相當(dāng)并且速度快了一倍，但只用了其27%的參數(shù)和20%的計算量。當(dāng)分辨率增加到1024時，ViG-T相較于基于Transformer的DeiT-T將計算量縮減5.2倍，GPU顯存節(jié)省90%，速度快了3.8倍，精度更是高了20.7%。這些結(jié)果證實了ViG作為一個高效且可縮放的基礎(chǔ)視覺骨干網(wǎng)絡(luò)的廣闊潛力。

下圖則是進(jìn)一步顯示了，我們隨著分辨率增大而凸顯的性能優(yōu)勢，即隨著輸入圖片分辨率的增大，ViG計算量更低，顯存占用更少，速度更快，精度更高。

路線圖也充分展示了設(shè)計的有效性。

我們的雙向設(shè)計只引入了2%的參數(shù)量便得到了11.1精度的提升，同時我們設(shè)計的硬件感知實現(xiàn)減少了19%的推理代價和13%的顯存占用，使其更加適合低算力場景的部署運(yùn)行。

總結(jié)和展望

本工作提出的 ViG 方法在高效視覺序列建模領(lǐng)域展現(xiàn)了重要的創(chuàng)新性和實際應(yīng)用價值。通過引入雙向門控線性注意力 (BiGLA) 和二維門控局部注入機(jī)制，ViG 有效彌補(bǔ)了現(xiàn)有方法在全局感知能力和局部細(xì)節(jié)捕獲上的不足，實現(xiàn)了兼具高效性與準(zhǔn)確性的視覺表示學(xué)習(xí)。這一方法不僅在圖像分類、目標(biāo)檢測和語義分割等任務(wù)中表現(xiàn)卓越，還顯著提升了模型在高分辨率場景下的性能與資源利用效率。

尤其值得注意的是，ViG 的硬件感知優(yōu)化設(shè)計充分結(jié)合了計算設(shè)備的硬件特性，大幅降低了內(nèi)存占用和計算成本。通過在多種任務(wù)中的驗證，ViG 展現(xiàn)了出色的精度、參數(shù)效率、顯存優(yōu)化和推理速度，為未來高清視覺信號處理以及基于大型語言模型 (LLM) 的多模態(tài)序列表征學(xué)習(xí)提供了強(qiáng)有力的技術(shù)支撐。