港大提出SparX:強(qiáng)化Vision Mamba和Transformer的稀疏跳躍連接機(jī)制

本文分享香港大學(xué)計(jì)算和數(shù)據(jù)科學(xué)學(xué)院俞益洲教授及其研究團(tuán)隊(duì)發(fā)表于 AAAI 2025 的論文——SparX，一種強(qiáng)化 Vision Mamba 和 Transformer 的稀疏跳躍連接機(jī)制，性能強(qiáng)大，代碼已開源。

論文標(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é)計(jì)算和數(shù)據(jù)科學(xué)學(xué)院俞益洲教授 (https://i.cs.hku.hk/~yzyu/index.html) 及其研究團(tuán)隊(duì)開發(fā)了一種新型的層間稀疏跳躍連接機(jī)制 --- 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é)細(xì)胞 (Retinal Ganglion Cell) 的啟發(fā), 將網(wǎng)絡(luò)的基礎(chǔ)模塊分為兩種不同的類型：神經(jīng)節(jié)層 (Ganglion Layer) 和常規(guī)層(Normal Layer)。

前者具有更高的復(fù)雜度和連接度，類似于視覺系統(tǒng)中的神經(jīng)節(jié)細(xì)胞，而后者的連接度和復(fù)雜度都較低，類似于視覺系統(tǒng)中的常規(guī)細(xì)胞。通過交叉堆疊 Ganglion Layer 和 Normal Layer 構(gòu)建了一種新的 Vision Backbone 網(wǎng)絡(luò)，在圖像分類、語義分割和目標(biāo)檢測中展現(xiàn)了強(qiáng)大的性能。

例如，基于 SparX 構(gòu)建的 Vision Mamba 模型 SparX-Mamba 相較于強(qiáng)大的 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)越性。

動機(jī)

隨著自注意力（Self-attention）和狀態(tài)空間模型（State Space Models 或 SSMs）在 NLP 任務(wù)中的優(yōu)異表現(xiàn)，許多工作將這些方法成功應(yīng)用到了視覺領(lǐng)域，例如 Swin-Transformer 和 VMamba。 目前主流的 Vision Backbone 模型的設(shè)計(jì)策略為構(gòu)建新的 token mixer，并據(jù)此來構(gòu)建視覺網(wǎng)絡(luò)。然而，不同層輸出的特征具有一定的互補(bǔ)性和冗余度，因此，尋找這些特征之間的互補(bǔ)性，并且移除冗余的特征可以進(jìn)一步提升網(wǎng)絡(luò)的表征能力，進(jìn)而提升性能。 雖然一些先前的工作（例如 DenseNet）已經(jīng)利用了不同網(wǎng)絡(luò)層的交互和復(fù)用來提升性能，但是 DenseNet 中的稠密連接具有較高的計(jì)算復(fù)雜度，使其難以直接用于比卷積更加復(fù)雜的和 SSM 算子。因此，設(shè)計(jì)一種高效的神經(jīng)連接機(jī)制來挖掘和利用網(wǎng)絡(luò)的層間互補(bǔ)性仍然需要進(jìn)一步探索。

方法

以基于 Mamba 的模型為例，在 SparX 中 Ganglion Layer 包含用于提取局部信息的動態(tài)位置編碼（Dynamic Position Encoding (DPE)），用于實(shí)現(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的密度，進(jìn)而控制網(wǎng)絡(luò)的復(fù)雜度和連接度，研究團(tuán)隊(duì)引入了一個(gè)步長參數(shù) S，即S=兩個(gè)最近的ganglion layers之間的normal layers的層數(shù)加一。

此外，該規(guī)則定義了兩種不同的連接類型：ganglion layer 和 norma layer 之間的內(nèi)連以及兩個(gè) ganglion layers 之間的互連。為了讓網(wǎng)絡(luò)具備強(qiáng)大的層間特征交互且具有高效性，ganglion layer 只與處于自身和最近的前一個(gè) ganglion layer 之間的那些 normal layers 建立內(nèi)連，但是同時(shí)與多個(gè)先前的 ganglion layers 建立互連。

這種設(shè)計(jì)是因?yàn)椋篻anglion layer 可以被視為網(wǎng)絡(luò)的“信息中心”，從最近的 normal layer 收集信息并與其他 ganglion layers 交換信息。一個(gè)簡單的例子為：如果一個(gè)網(wǎng)絡(luò)有 8 層并設(shè)置 S=2，則 normal layer 的索引為 {1, 3, 5, 7}，而 ganglion layers 的索引為 {2, 4, 6, 8}。

2. 跨層滑動窗口（Cross-layer Sliding Window）旨在進(jìn)一步提升網(wǎng)絡(luò)高效性，其設(shè)計(jì)背后的動機(jī)是：盡管上述層間連接方式具有稀疏性，深層網(wǎng)絡(luò)仍可能因需要存儲和訪問大量先前的特征圖而產(chǎn)生較高的顯存消耗。為此，受啟發(fā)于經(jīng)典的空間滑動窗口，引入另一個(gè)超參數(shù) M 來限制每個(gè) ganglion layer 僅與先前的 M 個(gè)最接近的 ganglion layers 建立互連。

基于這兩條新規(guī)則，即使沒有直接連接，語義信息仍然可以通過相對稀疏的內(nèi)連和互連從較淺的網(wǎng)絡(luò)層快速傳遞到較深的網(wǎng)絡(luò)層。圖 1 展示了一個(gè) SparX (S=2, M=2) 的示例。

▲圖1. Mamba-based SparX示例

為了選擇性地從先前網(wǎng)絡(luò)層的特征中挖掘有用的特征，從而動態(tài)的建模層間交互，研究團(tuán)隊(duì)提出了一個(gè)簡單有效的 DMCA 模塊。如下圖所示，DMCA 用當(dāng)前層特征作為 query，用先前層的特征作為 key/value 來構(gòu)建 channel-wise cross attention。構(gòu)建 channel attention 的目的是為了更好地進(jìn)行通道之間的信息交互，從而獲得更好的性能。

▲圖2. DMCA架構(gòu)圖

實(shí)驗(yàn)結(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 上也進(jìn)行了全面的評估，其性能在與一些強(qiáng)大的 Vision Backbones 的比較中脫穎而出，并且有著好的 tradeoff。

如表 3 所示，在 COCO 2017 數(shù)據(jù)集上，SparX 同樣展示出了更優(yōu)的性能。值得注意的是，當(dāng)使用更加強(qiáng)大的訓(xùn)練條件（3× schedule）時(shí)，SparX 展現(xiàn)了更加顯著的性能提升。

▲表3. COCO目標(biāo)檢測和實(shí)例分割性能對比

通用性實(shí)驗(yàn)

研究團(tuán)隊(duì)用經(jīng)典的 Swin-Transformer 測試了 SparX 的通用性，為了保持公平對比，micro 設(shè)計(jì)嚴(yán)格保持了和 Swin 原始設(shè)計(jì)一致，例如完全相同的 patch embedding 和 token mixer。如表 4 所示，SparX 在不同任務(wù)上均取得了顯著的性能提升。

▲表4. SparX用于Transformer架構(gòu)時(shí)的性能

消融實(shí)驗(yàn)

為了驗(yàn)證 SparX 的有效性，研究團(tuán)隊(duì)構(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的對比

可視化實(shí)驗(yàn)

Centered Kernel Alignment (CKA) 分析：不難發(fā)現(xiàn)，在表 5 消融實(shí)驗(yàn)中，連接度更高的模型（DGC 和 DSN）并沒有帶來性能提升。為了尋找背后的原因，我們進(jìn)行了模型的 CKA 分析。 如圖 3 所示，VMamba-T 每一層學(xué)習(xí)到的特征與相鄰層非常相似，說明了模型具有大量的特征冗余。此外，DGC 和 DSN 模型可以學(xué)習(xí)到更加多樣化的特征，降低了特征冗余度。與這些方法相比，SparX 不同層的特征更加多樣化，從而具有更加強(qiáng)大的特征表達(dá)，這也是其性能更好的原因。

▲圖3. CKA可視化

有效感受野（Effective Receptive Field 或 ERF）分析：如圖 4 所示，SparX-Mamba 和其它方法相比具有更大的感受野，進(jìn)一步說明了 SparX 對模型表征能力的增強(qiáng)效果。

▲圖4. Effective Receptive Fields可視化