ViG基于視覺Mamba的通用視覺主干網絡

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

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

? 論文鏈接：

https://arxiv.org/abs/2405.18425

? 本工作已被AAAI 2025收錄。

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

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

背景介紹

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

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

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

方法概述

雙向門控注意力機制

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

硬件感知的雙向算子設計

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

長短上下文的動態門控制機

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

結構設計

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

實驗結果

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

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

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

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

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

總結和展望

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

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