本文分享香港大學(xué)計算和數(shù)據(jù)科學(xué)學(xué)院俞益洲教授及其研究團隊發(fā)表于 AAAI 2025 的論文——SparX,一種強化 Vision Mamba 和 Transformer 的稀疏跳躍連接機制,性能強大,代碼已開源。
論文標(biāo)題: SparX: A Sparse Cross-Layer Connection Mechanism for Hierarchical Vision Mamba and Transformer Networks
論文鏈接:
https://arxiv.org/abs/2409.09649
代碼鏈接:
https://github.com/LMMMEng/SparX
摘要
香港大學(xué)計算和數(shù)據(jù)科學(xué)學(xué)院俞益洲教授 (https://i.cs.hku.hk/~yzyu/index.html) 及其研究團隊開發(fā)了一種新型的層間稀疏跳躍連接機制 --- Sparse Cross-Layer Connection Mechanism (SparX),可以有效提升包含 Vision Mamba 和 Transformer 在內(nèi)的 Vision Backbone 的性能。
不同于常規(guī)的 Vision Backbone 通過堆疊相同的基礎(chǔ)模塊來搭建網(wǎng)絡(luò)架構(gòu),SparX 受到了人類視覺系統(tǒng)中神經(jīng)節(jié)細胞 (Retinal Ganglion Cell) 的啟發(fā), 將網(wǎng)絡(luò)的基礎(chǔ)模塊分為兩種不同的類型:神經(jīng)節(jié)層 (Ganglion Layer) 和常規(guī)層(Normal Layer)。
前者具有更高的復(fù)雜度和連接度,類似于視覺系統(tǒng)中的神經(jīng)節(jié)細胞,而后者的連接度和復(fù)雜度都較低,類似于視覺系統(tǒng)中的常規(guī)細胞。通過交叉堆疊 Ganglion Layer 和 Normal Layer 構(gòu)建了一種新的 Vision Backbone 網(wǎng)絡(luò),在圖像分類、語義分割和目標(biāo)檢測中展現(xiàn)了強大的性能。
例如,基于 SparX 構(gòu)建的 Vision Mamba 模型 SparX-Mamba 相較于強大的 VMamba 仍有明顯提升:雖然參數(shù)量更少,SparX-Mamba-T 在 ImageNet-1K 上的 Top-1 準(zhǔn)確率仍超越 VMamba-T 1%。此外,SparX-Mamba 在語義分割和目標(biāo)檢測任務(wù)上同樣具備突出的性能,展現(xiàn)了 Mamba 模型在學(xué)習(xí)長距離關(guān)聯(lián)方面的優(yōu)越性。
動機
隨著自注意力(Self-attention)和狀態(tài)空間模型(State Space Models 或 SSMs)在 NLP 任務(wù)中的優(yōu)異表現(xiàn),許多工作將這些方法成功應(yīng)用到了視覺領(lǐng)域,例如 Swin-Transformer 和 VMamba。 目前主流的 Vision Backbone 模型的設(shè)計策略為構(gòu)建新的 token mixer,并據(jù)此來構(gòu)建視覺網(wǎng)絡(luò)。然而,不同層輸出的特征具有一定的互補性和冗余度,因此,尋找這些特征之間的互補性,并且移除冗余的特征可以進一步提升網(wǎng)絡(luò)的表征能力,進而提升性能。 雖然一些先前的工作(例如 DenseNet)已經(jīng)利用了不同網(wǎng)絡(luò)層的交互和復(fù)用來提升性能,但是 DenseNet 中的稠密連接具有較高的計算復(fù)雜度,使其難以直接用于比卷積更加復(fù)雜的和 SSM 算子。因此,設(shè)計一種高效的神經(jīng)連接機制來挖掘和利用網(wǎng)絡(luò)的層間互補性仍然需要進一步探索。
方法
以基于 Mamba 的模型為例,在 SparX 中 Ganglion Layer 包含用于提取局部信息的動態(tài)位置編碼(Dynamic Position Encoding (DPE)),用于實現(xiàn)層間信息交互的 Dynamic Multi-layer Channel Aggregator (DMCA),和用于空間上下文信息建模的視覺狀態(tài)空間模型(Visual State Space Model 或 VSS),而 Normal Layer 則沒有 DMCA 模塊。
此外,SparX 引入了兩條新的跨層連接規(guī)則:
1. Sparse Ganglion Layers 將一組具有均勻間隔的層指定為更加復(fù)雜且連接度更高的 ganglion layers,而所有其余層則為復(fù)雜度和連接度都較低的 normal layers。為了控制ganglion layers的密度,進而控制網(wǎng)絡(luò)的復(fù)雜度和連接度,研究團隊引入了一個步長參數(shù) S,即S=兩個最近的ganglion layers之間的normal layers的層數(shù)加一。
此外,該規(guī)則定義了兩種不同的連接類型:ganglion layer 和 norma layer 之間的內(nèi)連以及兩個 ganglion layers 之間的互連。為了讓網(wǎng)絡(luò)具備強大的層間特征交互且具有高效性,ganglion layer 只與處于自身和最近的前一個 ganglion layer 之間的那些 normal layers 建立內(nèi)連,但是同時與多個先前的 ganglion layers 建立互連。
這種設(shè)計是因為:ganglion layer 可以被視為網(wǎng)絡(luò)的“信息中心”,從最近的 normal layer 收集信息并與其他 ganglion layers 交換信息。一個簡單的例子為:如果一個網(wǎng)絡(luò)有 8 層并設(shè)置 S=2,則 normal layer 的索引為 {1, 3, 5, 7},而 ganglion layers 的索引為 {2, 4, 6, 8}。
2. 跨層滑動窗口(Cross-layer Sliding Window)旨在進一步提升網(wǎng)絡(luò)高效性,其設(shè)計背后的動機是:盡管上述層間連接方式具有稀疏性,深層網(wǎng)絡(luò)仍可能因需要存儲和訪問大量先前的特征圖而產(chǎn)生較高的顯存消耗。為此,受啟發(fā)于經(jīng)典的空間滑動窗口,引入另一個超參數(shù) M 來限制每個 ganglion layer 僅與先前的 M 個最接近的 ganglion layers 建立互連。
基于這兩條新規(guī)則,即使沒有直接連接,語義信息仍然可以通過相對稀疏的內(nèi)連和互連從較淺的網(wǎng)絡(luò)層快速傳遞到較深的網(wǎng)絡(luò)層。圖 1 展示了一個 SparX (S=2, M=2) 的示例。
▲圖1. Mamba-based SparX示例
為了選擇性地從先前網(wǎng)絡(luò)層的特征中挖掘有用的特征,從而動態(tài)的建模層間交互,研究團隊提出了一個簡單有效的 DMCA 模塊。如下圖所示,DMCA 用當(dāng)前層特征作為 query,用先前層的特征作為 key/value 來構(gòu)建 channel-wise cross attention。構(gòu)建 channel attention 的目的是為了更好地進行通道之間的信息交互,從而獲得更好的性能。
▲圖2. DMCA架構(gòu)圖
實驗結(jié)果
圖像分類:SparX 在大規(guī)模數(shù)據(jù)集 ImageNet-1K 上表現(xiàn)出了卓越的性能,相較于現(xiàn)有方法,展現(xiàn)出更為出色的性能以及更好的 tradeoff。如表 1 所示,SparX-Mamba-S 僅以大約一半的參數(shù)量和 FLOPs 就超越了 VMamba-B 的性能。
如表 2 所示,對 SparX 在 ADE20K 上也進行了全面的評估,其性能在與一些強大的 Vision Backbones 的比較中脫穎而出,并且有著好的 tradeoff。
如表 3 所示,在 COCO 2017 數(shù)據(jù)集上,SparX 同樣展示出了更優(yōu)的性能。值得注意的是,當(dāng)使用更加強大的訓(xùn)練條件(3× schedule)時,SparX 展現(xiàn)了更加顯著的性能提升。
▲表3. COCO目標(biāo)檢測和實例分割性能對比
通用性實驗
研究團隊用經(jīng)典的 Swin-Transformer 測試了 SparX 的通用性,為了保持公平對比,micro 設(shè)計嚴(yán)格保持了和 Swin 原始設(shè)計一致,例如完全相同的 patch embedding 和 token mixer。如表 4 所示,SparX 在不同任務(wù)上均取得了顯著的性能提升。
▲表4. SparX用于Transformer架構(gòu)時的性能
消融實驗
為了驗證 SparX 的有效性,研究團隊構(gòu)建了兩種不同的稠密連接模型:1)Dense Ganglion Connections (DGC-Mamba-T):移除跨層滑動窗口(Cross-layer Sliding Window)來消除互連稀疏性;2)DenseNet-style Network (DSN-Mamba-T):完全按照 DenseNet 的策略來構(gòu)建模型。
如表 5 所示,SparX 在保持最優(yōu)性能的前提下還具備高效性。
▲表5. SparX和Dense Connection的對比
可視化實驗
Centered Kernel Alignment (CKA) 分析:不難發(fā)現(xiàn),在表 5 消融實驗中,連接度更高的模型(DGC 和 DSN)并沒有帶來性能提升。為了尋找背后的原因,我們進行了模型的 CKA 分析。 如圖 3 所示,VMamba-T 每一層學(xué)習(xí)到的特征與相鄰層非常相似,說明了模型具有大量的特征冗余。此外,DGC 和 DSN 模型可以學(xué)習(xí)到更加多樣化的特征,降低了特征冗余度。與這些方法相比,SparX 不同層的特征更加多樣化,從而具有更加強大的特征表達,這也是其性能更好的原因。
▲圖3. CKA可視化
有效感受野(Effective Receptive Field 或 ERF)分析:如圖 4 所示,SparX-Mamba 和其它方法相比具有更大的感受野,進一步說明了 SparX 對模型表征能力的增強效果。
▲圖4. Effective Receptive Fields可視化
-
開源
+關(guān)注
關(guān)注
3文章
3547瀏覽量
43351 -
連接
+關(guān)注
關(guān)注
2文章
101瀏覽量
21211 -
Vision
+關(guān)注
關(guān)注
1文章
201瀏覽量
18595
原文標(biāo)題:AAAI 2025 | 港大提出SparX:強化Vision Mamba和Transformer的稀疏跳躍連接機制
文章出處:【微信號:tyutcsplab,微信公眾號:智能感知與物聯(lián)網(wǎng)技術(shù)研究所】歡迎添加關(guān)注!文章轉(zhuǎn)載請注明出處。
發(fā)布評論請先 登錄
首個Mamba系列模型量化方案MambaQuant解讀

字節(jié)豆包大模型團隊提出稀疏模型架構(gòu)
如何使用MATLAB構(gòu)建Transformer模型

SparX-5/5i硬件設(shè)計檢查清單

ADS805E測量跳躍幅度很大,不知道是為什么?
王欣然教授團隊提出基于二維材料的高效稀疏神經(jīng)網(wǎng)絡(luò)硬件方案

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

一種新的通用視覺主干模型Vision Mamba

Mamba入局圖像復(fù)原,達成新SOTA

HTC新頭顯Vive Focus Vision:強化混合現(xiàn)實體驗
港迪技術(shù)IPO上市丨提升綜合管理水平,強化公司持續(xù)經(jīng)營能力
Transformer能代替圖神經(jīng)網(wǎng)絡(luò)嗎
Transformer語言模型簡介與實現(xiàn)過程
Transformer 能代替圖神經(jīng)網(wǎng)絡(luò)嗎?

評論