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

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

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

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

? 論文鏈接：

https://arxiv.org/abs/2405.18425

? 本工作已被AAAI 2025收錄。

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

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

背景介紹

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

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

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

方法概述

雙向門控注意力機制

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

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

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

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

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

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

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

實驗結(jié)果

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

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

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

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

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

總結(jié)和展望

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

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