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

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

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

為了進(jìn)一步推進(jìn)線性復(fù)雜度視覺序列表征學(xué)習(xí)的效率，我們結(jié)合當(dāng)今時(shí)代計(jì)算設(shè)備的硬件特性去設(shè)計(jì)結(jié)構(gòu)，將自然語(yǔ)言序列建模中硬件計(jì)算更為友好的門控線性注意力模塊GLA引入到視覺表征學(xué)習(xí)中，進(jìn)一步設(shè)計(jì)參數(shù)高效的雙向建模、長(zhǎng)短上下文的動(dòng)態(tài)門控機(jī)制和硬件感知的雙向算子設(shè)計(jì)。通過以上設(shè)計(jì)，我們提出的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)勢(shì)，但在更為常見的分辨率，仍略遜于Transformer和CNN。

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

背景介紹

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

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

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

方法概述

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

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

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

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

長(zhǎng)短上下文的動(dòng)態(tài)門控制機(jī)

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

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

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

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

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

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

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

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

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

總結(jié)和展望

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

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