52CV 2024-12-24 21:21 江苏

性能强大，代码开源!

关注公众号，发现CV技术之美

本篇分享 AAAI 2025 论文SparX: A Sparse Cross-Layer Connection Mechanism for Hierarchical Vision Mamba and Transformer Networks，香港大学提出 SparX，一种强化 Vision Mamba 和 Transformer 的稀疏跳跃连接机制，性能强大，代码已开源!

论文链接：https://arxiv.org/abs/2409.09649

代码链接：https://github.com/LMMMEng/SparX

摘要

香港大学计算和数据科学学院俞益洲教授(https://i.cs.hku.hk/~yzyu/index.html)及其研究团队开发了一种新型的层间稀疏跳跃连接机制 --- Sparse Cross-Layer Connection Mechanism (SparX)，可以有效提升包含Vision Mamba和Transformer在内的Vision Backbone的性能。

不同于常规的Vision Backbone通过堆叠相同的基础模块来搭建网络架构，SparX受到了人类视觉系统中神经节细胞 (Retinal Ganglion Cell) 的启发, 将网络的基础模块分为两种不同的类型：神经节层(Ganglion Layer)和常规层(Normal Layer)。

前者具有更高的复杂度和连接度，类似于视觉系统中的神经节细胞，而后者的连接度和复杂度都较低，类似于视觉系统中的常规细胞。通过交叉堆叠Ganglion Layer和Normal Layer构建了一种新的Vision Backbone网络，在图像分类、语义分割和目标检测中展现了强大的性能。

例如，基于SparX构建的Vision Mamba模型SparX-Mamba相较于强大的VMamba仍有明显提升：虽然参数量更少，SparX-Mamba-T在ImageNet-1K上的Top-1准确率仍超越VMamba-T 1%。此外，SparX-Mamba在语义分割和目标检测任务上同样具备突出的性能，展现了Mamba模型在学习长距离关联方面的优越性。

动机

随着自注意力（Self-attention）和状态空间模型（State Space Models或SSMs)在NLP任务中的优异表现，许多工作将这些方法成功应用到了视觉领域，例如Swin-Transformer和VMamba。

目前主流的Vision Backbone模型的设计策略为构建新的token mixer，并据此来构建视觉网络。然而，不同层输出的特征具有一定的互补性和冗余度，因此，寻找这些特征之间的互补性，并且移除冗余的特征可以进一步提升网络的表征能力，进而提升性能。

虽然一些先前的工作(例如DenseNet)已经利用了不同网络层的交互和复用来提升性能，但是DenseNet中的稠密连接具有较高的计算复杂度，使其难以直接用于比卷积更加复杂的和SSM算子。因此，设计一种高效的神经连接机制来挖掘和利用网络的层间互补性仍然需要进一步探索。

方法

以基于Mamba的模型为例，在SparX中Ganglion Layer包含用于提取局部信息的动态位置编码（Dynamic Position Encoding (DPE)），用于实现层间信息交互的Dynamic Multi-layer Channel Aggregator (DMCA)，和用于空间上下文信息建模的视觉状态空间模型（Visual State Space Model 或VSS)，而Normal Layer则没有DMCA模块。此外，SparX引入了两条新的跨层连接规则：

Sparse Ganglion Layers将一组具有均匀间隔的层指定为更加复杂且连接度更高的ganglion layers，而所有其余层则为复杂度和连接度都较低的normal layers。为了控制ganglion layers的密度，进而控制网络的复杂度和连接度，研究团队引入了一个步长参数S，即S=两个最近的ganglion layers之间的normal layers的层数加一。此外，该规则定义了两种不同的连接类型: ganglion layer和norma layer之间的内连以及两个ganglion layers之间的互连。为了让网络具备强大的层间特征交互且具有高效性，ganglion layer只与处于自身和最近的前一个ganglion layer之间的那些normal layers建立内连，但是同时与多个先前的ganglion layers建立互连。这种设计是因为：ganglion layer可以被视为网络的“信息中心”，从最近的normal layer收集信息并与其他ganglion layers交换信息。一个简单的例子为：如果一个网络有8层并设置S=2，则normal layer的索引为{1, 3, 5, 7}，而ganglion layers的索引为{2, 4, 6, 8}。

跨层滑动窗口（Cross-layer Sliding Window）旨在进一步提升网络高效性，其设计背后的动机是：尽管上述层间连接方式具有稀疏性，深层网络仍可能因需要存储和访问大量先前的特征图而产生较高的显存消耗。为此，受启发于经典的空间滑动窗口，引入另一个超参数M来限制每个ganglion layer仅与先前的M个最接近的ganglion layers建立互连。基于这两条新规则，即使没有直接连接，语义信息仍然可以通过相对稀疏的内连和互连从较浅的网络层快速传递到较深的网络层。图1展示了一个SparX (S=2, M=2) 的示例。

为了选择性地从先前网络层的特征中挖掘有用的特征，从而动态的建模层间交互，研究团队提出了一个简单有效的DMCA模块。如下图所示，DMCA用当前层特征作为query，用先前层的特征作为key/value来构建channel-wise cross attention。构建channel attention的目的是为了更好地进行通道之间的信息交互，从而获得更好的性能。

实验结果

图像分类：SparX在大规模数据集 ImageNet-1K 上表现出了卓越的性能，相较于现有方法，展现出更为出色的性能以及更好的tradeoff。如表1所示，SparX-Mamba-S仅以大约一半的参数量和FLOPs就超越了VMamba-B的性能。

如表 2 所示，对SparX在ADE20K上也进行了全面的评估，其性能在与一些强大的 Vision Backbones 的比较中脱颖而出，并且有着好的tradeoff。

如表3所示，在 COCO 2017 数据集上，SparX同样展示出了更优的性能。值得注意的是，当使用更加强大的训练条件（3× schedule）时，SparX展现了更加显著的性能提升。

通用性实验

研究团队用经典的Swin-Transformer测试了SparX的通用性，为了保持公平对比，micro设计严格保持了和Swin原始设计一致，例如完全相同的patch embedding和token mixer。如表4所示，SparX在不同任务上均取得了显著的性能提升。

消融实验

为了验证SparX的有效性，研究团队构建了两种不同的稠密连接模型：(1) Dense Ganglion Connections (DGC-Mamba-T)：移除跨层滑动窗口（Cross-layer Sliding Window）来消除互连稀疏性；(2) DenseNet-style Network (DSN-Mamba-T): 完全按照DenseNet的策略来构建模型。如表5所示，SparX在保持最优性能的前提下还具备高效性。

可视化实验

Centered Kernel Alignment (CKA)分析：不难发现，在表5消融实验中，连接度更高的模型(DGC和DSN)并没有带来性能提升。为了寻找背后的原因，我们进行了模型的CKA分析。如图3所示，VMamba-T 每一层学习到的特征与相邻层非常相似，说明了模型具有大量的特征冗余。此外，DGC 和 DSN 模型可以学习到更加多样化的特征，降低了特征冗余度。与这些方法相比，SparX不同层的特征更加多样化，从而具有更加强大的特征表达，这也是其性能更好的原因。

有效感受野（Effective Receptive Field 或ERF)分析：如图4所示，SparX-Mamba和其它方法相比具有更大的感受野，进一步说明了SparX对模型表征能力的增强效果。