SparX：一种用于层次视觉Mamba和变换器网络的稀疏跨层连接机制

摘要

https://arxiv.org/pdf/2409.09649
由于动态状态空间模型（SSMs）能够以线性时间计算复杂度捕获长距离依赖关系，Mamba在自然语言处理（NLP）任务中表现出了显著的性能。这激发了基于Mamba的视觉模型的快速发展，从而在视觉识别任务中取得了令人瞩目的成果。然而，这类模型不能通过特征聚合、交互和选择来跨层提炼特征。此外，为卷积神经网络（CNNs）或视觉Transformer（ViTs）设计的现有跨层特征聚合方法在基于Mamba的模型中由于计算成本高而不切实际。因此，本文旨在为视觉主干网络引入一种高效的跨层特征聚合机制。受人类视觉系统中视网膜神经节细胞（RGCs）的启发，我们提出了一种新的稀疏跨层连接机制，称为SparX，以有效改进跨层特征交互和复用。具体来说，我们构建了两种不同类型的网络层：神经节层和普通层。前者具有更高的连接性和复杂性，能够以输入依赖的方式实现多层特征聚合和交互。相比之下，后者的连接性和复杂性较低。通过交错这两种类型的层，我们设计了一个新的视觉主干网络系列，该系列具有稀疏跨连接层，与同类网络相比，在模型大小、计算成本、内存成本和准确性方面实现了良好的权衡。例如，在参数更少的情况下，SparX-Mamba-T将VMambaT的top-1准确率从$82.5\%$提高到$83.5\%$，而SparX-Swin-T与Swin-T相比，top-1准确率提高了$1.3\%$。广泛的实验结果表明，我们的新连接机制在各种视觉任务上具有优越的性能和泛化能力。代码公开可用，网址为https://github.com/LMMMEng/SparX。

1 引言

上下文建模在计算机视觉任务中发挥着至关重要的作用，包括图像分类、目标检测和语义分割。最近的研究探索了大卷积核以扩大感受野（Liu等人，2022；Ding等人，2022b；Liu等人，2023；Ding等人，2024）以及视觉Transformer（ViTs）（Dosovitskiy等人，2021；Touvron等人，2021；Liu等人，2021；Wang等人，2021），使其具备强大的长距离建模能力。然而，无论是具有大卷积核的CNN还是视觉Transformer都存在局限性。CNN的性能仍然落后于先进的ViT，而后者在处理高分辨率输入时由于二次复杂度而面临高昂的计算成本。

凭借能够以接近线性复杂度建模长距离依赖关系的动态状态空间模型（SSMs），Mamba（Gu和Dao，2024）在自然语言处理（NLP）任务中表现出了良好的性能。基于这一成功，研究人员提出了用于视觉任务的基于Mamba的模型（Zhu等人，2024；Yang等人，2024；Liu等人，2024；Huang等人，2024；Behrouz, Santacatterina和Zabih，2024；Pei, Huang和Xu，2024）。这类工作通过利用SSMs（主要用于捕获同一层内的空间上下文信息）成功地将线性时间长距离建模能力引入视觉模型。然而，不同层的标记特征之间可能存在互补性和冗余性，因为它们在不同层次和粒度上捕获图像的特征和语义。因此，需要通过聚合、交互和选择来跨层提炼特征，以提取有用信息并去除冗余。然而，现有的基于Mamba的视觉模型并不具备跨层特征提炼的能力，这阻碍了SSMs全部潜力的发挥。

尽管跨层特征交互和复用已被证明在CNN和视觉Transformer中都能有效提高性能，如DenseNet（Huang等人，2017）和FcaFormer（Zhang, Hu和Wang，2023），但这些方法并不直接适用于基于Mamba的视觉模型。如表5所示，当将VMamba重构为DenseNet风格的网络（DSN）时，其吞吐量下降了近$50\%$，同时GPU内存使用量增加了1 GB。与此同时，如表6所示，在FcaFormer中使用空间交叉注意力（SCA）作为多层特征交互模块，与简单的特征拼接相比，由于注意力计算的二次复杂度，导致GPU内存使用量增加了$80\%$以上。因此，在基于Mamba的模型中设计一个高效的连接模式来促进跨层特征交互是一项重大挑战。

与主要关注改进基于Mamba的令牌混合器（token mixer）的现有工作不同，在本工作中，我们提出了一种新的架构设计，该设计显著提升了基于SSM（Space-to-Space Mapping，空间到空间映射）的视觉模型的性能。我们的设计灵感来源于人类视觉系统中的视网膜神经节细胞（Retinal Ganglion Cells，RGCs）。RGCs在人类视觉系统中作为信息传输中心，连接视网膜输入和中枢神经系统中的视觉处理中心（Kim等人，2021）。由于RGC层内复杂的神经结构和多样的细胞组成（Curcio等人，1990；Watson，2014），它与其他视觉系统中的细胞层有大量的连接，从而与其他层产生复杂的神经交互。受这一生物机制的启发，我们提出了一种新的稀疏跨层连接机制，命名为SparX。我们定义了两种基于SSM的基本层，即普通层和神经节层。前者信息流较少，从前一层获取单个输入并与少数后续层连接。相比之下，后者与后续层建立了更多的连接，并对来自多个前层的输入进行编码。神经节层与普通层的组合模仿了人类视觉系统中RGC层和非RGC层的组合。

现有文献（D’Souza和Lang，2020）也证实，RGCs产生多种细胞间通信，包括RGC-RGC和RGC-非RGC相互作用。受这一前提的启发，我们在神经节层中引入了一种新的动态多层通道聚合器（Dynamic Multi-layer Channel Aggregator，DMCA），旨在有效地促进多层之间的自适应特征聚合和交互。此外，为了保持速度和降低内存消耗，我们引入了一个跨层滑动窗口，它只允许每个神经节层与同一滑动窗口中的其他神经节层连接。

通过在前馈网络中分层交错排列神经节层和普通层，我们设计了一种新的基于Mamba的视觉主干网络，即具有稀疏跨连接层的Vision Mamba（SparX-Mamba）。除了基于Mamba的模型的固有优势外，我们的SparX-Mamba在性能和计算成本之间实现了良好的权衡。值得注意的是，我们的SparX-Mamba不同于通过简单堆叠相同层设计而构建的传统视觉模型，导致每一层具有相似的特征提取能力。相比之下，SparX-Mamba的两种功能不同的层有助于模型提取更多样化的特征表示，从而提高了性能。值得注意的是，我们的SparX还已成功应用于分层视觉Transformer，如Swin（Liu等人，2021），并且所得模型在原始Transformer模型的基础上也具有上述优势。例如，与Swin-T相比，我们的SparX-Swin-T在top-1准确率上提高了1.3%。

如图1所示，所提出的SparX-Mamba在最近提出的VMamba以及经典的基于卷积神经网络（CNN）和Transformer的模型上表现出优越的性能，且参数数量（Params）更少。例如，与VMamba-T相比，SparX-MambaT在ImageNet-1K上实现了83.5%的top-1准确率，显著超过了VMamba-T的82.5%准确率。当与UperNet结合用于ADE20K数据集上的语义分割时，SparX-Mamba-T在平均交并比（mean Intersection over Union，mIoU）上比VMamba-T高出1.7%。此外，当与Mask R-CNN结合用于COCO 2017数据集上的目标检测时，SparX-Mamba-T在平均精度（Average Precision，${\mathrm{AP}}^b$）上比VMamba-T高出1.3%，甚至超过了复杂性几乎是SparX-Mamba-T的1.5倍的VMamba-S。同时，SparX-Mamba的小型和基型版本也在性能上相对于VMamba有显著的提升。

综上所述，我们的主要贡献有三方面：首先，受人类视觉系统中RGCs的启发，我们提出了一种新的跳跃连接机制SparX，该机制稀疏且动态地配置跨层连接，从而实现多样化的信息流和改进的特征蒸馏。其次，在Mamba和Transformer的基础上，我们提出了两种通用的视觉主干网络SparX-Mamba和SparX-Swin，它们由两种功能不同的层构建，其中一层负责相对较远层之间的通信。第三，我们在图像分类、目标检测和语义分割任务上进行了广泛的实验。结果表明，我们的新网络架构在性能和计算成本之间实现了显著的权衡。

2 相关工作

卷积神经网络（CNNs）。卷积神经网络已成为计算机视觉中的标准深度学习架构（Simonyan和Zisserman，2015；He等，2016；Huang等，2017）。在现代CNN领域，已经从传统的小型卷积核显著转变为设计关注大型卷积核。这一趋势的显著例子包括ConvNeXt（Liu等，2022）、RepLKNet（Ding等，2022b，2024）和SLaK（Liu等，2023）。最近，InceptionNeXt（Yu等，2024b）将经典的Inception网络与ConvNeXt相结合，实现了显著的性能提升。
视觉Transformer。为适应NLP任务中的Transformer到计算机视觉任务，视觉Transformer（ViT）（Dosovitskiy等，2021）通过补丁嵌入将图像分割为视觉标记，从而允许多头自注意力（MHSA）学习标记到标记的依赖性。为进一步生成层次化特征表示并提高模型效率，许多后续工作采用了金字塔架构设计，包括窗口注意力（Liu等，2021；Dong等，2022；Pan等，2023）、稀疏注意力（Yang等，2021；Wang等，2022；Ren等，2022；Wu等，2023）以及卷积注意力混合模型（Dai等，2021；Xiao等，2021；Li等，2022，2023；Lou等，2023；Yu等，2024a）。

Vision Mamba。由于Mamba（Gu和Dao，2024）在NLP任务中取得了卓越性能，许多研究人员已将Mamba转移到计算机视觉任务。作为Mamba的核心，SSM能够以接近线性的复杂度建模长距离依赖，并在视觉任务中表现出色。例如，ViM（Zhu等，2024）引入了一个双向SSM模块，并构建了类似ViT（Dosovitskiy等，2021）的等距架构。同样，PlainMamba（Yang等，2024）也通过连续的二维扫描构建了等距架构。VMamba（Liu等，2024）将扫描顺序扩展到四个方向，是早期基于SSM的层次化架构。随后，一系列层次化视觉Mamba模型被提出，包括MambaMixer（Behrouz，Santacatterina和Zabih，2024）、LocalMamba（Huang等，2024）、EfficientVMamba（Pei，Huang和Xu，2024）以及MSVMamba（Shi，Dong和Xu，2024）。

跳跃连接。ResNet（He等，2016）为深度CNN引入了残差连接，通过绕过捷径缓解了梯度消失和梯度爆炸问题。为丰富信息流，DenseNet（Huang等，2017）进一步引入了密集连接，它使用所有前面层的输出集合来计算一层的输入。DPN（Chen等，2017）利用ResNet和DenseNet中的跳跃连接来构建双路径网络。在密集预测任务领域，U-Net（Ronneberger，Fischer和Brox，2015）和FPN（Lin等，2017）利用跳跃连接来弥合编码器中的低级细节与解码器中的高级上下文之间的鸿沟。本文提出了一种新的跳跃连接机制，该机制动态配置稀疏的跨层连接，与现有方法相比性能有所提升。

跨特征注意力。跨特征注意力可以通过改善不同特征之间的交互来增强特征表示。例如，许多工作（Chen，Fan和Panda，2021；Lee等，2022；Wang等，2023a；Ren等，2022；Wu等，2023）提出通过生成多尺度标记来执行跨特征注意力，然后将这些标记输入到自注意力或交叉注意力层以建模多尺度交互。最近，FcaFormer（Zhang，Hu和Wang，2023）引入了一层空间标记与前几层代表性标记之间的交叉注意力来建模跨层特征交互。与现有方法相比，本文引入了一种新的通道级跨注意力机制，该机制动态整合了来自前面层的特征通道，从而促进了高效的跨层特征交互。

3 方法

3.1 稀疏跨层连接

概述。受人类视觉系统中视网膜神经节细胞（RGCs）功能的启发，本文引入了一种名为SparX的新型稀疏跨层连接机制。其目标是有效地建模跨层通信，生成多样化的信息流，并改善基于Mamba架构中的特征重用。如图2所示，SparX有三个构建组件：动态位置编码（DPE）（Chu等，2022）、Mamba块和新引入的动态多层通道聚合器（DMCA）。其中，DPE利用了残差$3\times 3$深度卷积（DWConv），这在现代视觉骨干网络中得到了广泛应用（Li等，2023；Chu等，2021；Guo等，2022）。Mamba块以接近线性的计算复杂度有效地捕获长距离依赖。然而，原始Mamba块（Gu和Dao，2024）不适合直接嵌入到视觉骨干网络中。在这方面，本文使用了在VMamba（Liu等，2024）中引入后表现出良好性能的VSS块。为增强通道混合，在VSS之后按照（Shi，Dong和Xu，2024）融入了ConvFFN（Wang等，2022）。所提出的DMCA动态聚合、选择和编码跨层特征，从而产生强大且稳健的特征表示。在这些基本组件的基础上，SparX中构建了两种类型的网络层：神经节层（“DPE $\to$ DMCA $\to$ VSS块”）和普通层（“DPE $\to$ VSS块”）。这两种层在前馈网络中交替出现，能够以稀疏的方式有效地建模跨层特征交互。

连接规则。在原始密集连接（Huang等，2017）中，每一层都与所有前面层相连。尽管DenseNet的浮点运算数（FLOPs）相对较低，但需要在前向传播过程中反复访问的许多前面特征图的存储，导致内存成本高、速度慢。尽管已有一些工作（Pleiss等，2017；Huang等，2019）试图降低内存成本，但当需要将模型扩展到更深更宽的结构时，密集连接仍然面临挑战。因此，原始密集连接机制在计算上过于繁重，无法用于其标记混合器比标准卷积更复杂的基于Mamba的模型。为此，本文提出了SparX，它可以有效地集成到最近的基于Mamba的架构中。在SparX中，本文引入了两条新规则，包括神经节层的稀疏放置和使用跨层滑动窗口来控制跨层连接的密度。
稀疏节层（Sparse Ganglion Layers）旨在指定一组均匀间隔的层作为节层，而所有剩余的层均为普通层。为了控制节层的密度，我们定义了一个超参数，称为步长（$S$），即两个最近的节层之间的普通层数量加一。我们进一步定义了两种跨层连接：（1）节层与普通层之间的内部连接；（2）两个节层之间的互联。为了建立稀疏的跨层连接，一个节层仅与最近的前一个节层与其自身之间的普通层有内部连接，同时它与多个前面的节层建立互联。这种跨层连接的原理是，节层可以视为一个信息交换中心，从最近的普通层收集信息，并与其他节层交换信息。为了确保网络或阶段的最终输出包含丰富的语义信息，网络或阶段的最后一层通常是节层。例如，如果我们有一个8层的网络，并将步长（$S$）设为2，那么普通层的索引为 { 1 , 3 , 5 , 7 } \{1,3,5,7\} {1,3,5,7}，而节层的索引为 { 2 , 4 , 6 , 8 } \{2,4,6,8\} {2,4,6,8}。

受空间滑动窗口的启发，我们提出了跨层滑动窗口（Cross-layer Sliding Window），以进一步提高计算效率。这一设计的动机是，尽管使用了上述稀疏连接，但更深的网络仍可能因需要存储和访问大量早期的特征图而产生较高的内存成本。为此，我们引入了另一个超参数$\mathbf{M}$，它限制了一个节层仅与最近的$M$个前面的节层有互联。基于这两个新规则，即使没有直接连接，信息仍然可以通过相对较少的内部连接和互联从较浅的层快速流向较深的层。图2展示了SparX在$\mathbf{S}=2$和$\mathbf{M}=2$时的一个示例。

动态多层通道聚合器（Dynamic Multi-layer Channel Aggregator，DMCA）。为了从前面的层中选择性地检索互补特征并动态地建模多层交互，我们提出了一种高效的动态多层通道聚合器（DMCA）。如图3所示，设$\mathbf{X}\in {\mathbb{R}}^{C\times N}$表示一个节层的特征，而 { Y 1 , Y 2 , ⋯ , Y L } ∈ R C × N \left\{\mathbf{Y}_{\mathbf{1}}, \mathbf{Y}_{\mathbf{2}}, \cdots, \mathbf{Y}_{\mathbf{L}}\right\} \in \mathbb{R}^{C \times N} {Y1​,Y2​,⋯,YL​}∈RC×N表示来自前面的节层和普通层的特征，其中$C$和$N$分别表示通道和空间维度，而$\mathbf{L}$表示来自先前层的特征总数。首先，我们将来自前面层的特征进行拼接，并使用一个线性层将通道维度投影到$2C$。然后，我们将投影后的特征沿着通道维度均匀地分成两段，即${\mathbf{Y}}_{\mathbf{k}}$和${\mathbf{Y}}_{\mathbf{v}}$。之后，我们设计了一个跨层分组通道注意力（Cross-layer Grouped Channel Attention，CGCA）机制，其中查询（query）、键（key）和值（value）分别对应于$\mathbf{X}$、${\mathbf{Y}}_{\mathbf{k}}$和${\mathbf{Y}}_{\mathbf{v}}$。具体来说，我们将通道分成多个组，并在每个组内计算通道级的交叉注意力。在我们的实验中，组的数量（$G$）始终设置为4。因此，在我们的CGCA中，所有组的整体注意力图$\mathbf{A}$的大小为$G\times \frac{C}{G}\times \frac{C}{G}$，这在使用大的空间分辨率时避免了高的计算和内存开销。

另一方面，通道级的交叉注意力计算一个注意力矩阵，该矩阵衡量来自查询和键的每一对变换通道之间的相似性，并消除了空间维度。注意，空间分辨率直接影响计算这个注意力矩阵的成本，因为它决定了查询和键中每个通道的维度。一个关键的观察结果是，注意力矩阵的大小与查询和键的空间维度无关，因此可以使用它们的空间缩减版本来计算原始注意力矩阵的近似值。基于此，我们采用了一种空间缩减策略来提高计算效率。具体来说，我们使用空间缩减器将查询和键中的空间标记数量从$N$压缩到$\frac{N}{r}$。在我们的实验中，设置$r$的规则是使$\frac{N}{r}$等于网络最后阶段中的标记数量，遵循PVT（Wang等人，2021，2022）。尽管此操作可能与PVT中使用的空间缩减相似，但其目的不同。在PVT中，空间缩减被用于计算空间标记到区域的注意力，其中缩减应用于键和值。在我们的情况下，我们使用空间缩减来高效地计算通道到通道的注意力，其中缩减应用于查询和键。一旦计算出注意力矩阵，就将其用于聚合值的通道，同时保持原始的空间分辨率。来自通道级交叉注意力的最终特征$\mathbf{Z}$可以视为在查询的指导下选择的值特征。最后，我们将$\mathbf{X}$、${\mathbf{Y}}_{\mathrm{v}}$和$\mathbf{Z}$进行拼接，并通过一个线性层将通道数量减少到$2C$。数学上，DMCA的操作总结如下：

$\begin{array}{c}{\mathbf{Y}}_{\mathbf{k}},{\mathbf{Y}}_{\mathbf{v}}=\mathrm{Split}\left(W_1\left(\mathrm{Cat}\left({\mathbf{Y}}_1,{\mathbf{Y}}_1,\cdots ,{\mathbf{Y}}_{\mathbf{L}}\right)\right)\right)\\ \mathrm{Q}=W_2\left(R_1(\mathbf{X})\right),\mathrm{K}=W_3\left(R_2\left({\mathbf{Y}}_{\mathbf{k}}\right)\right),\mathrm{V}=W_4\left({\mathbf{Y}}_{\mathbf{v}}\right)\\ \mathbf{Z}=\frac{\mathrm{Softmax}\left({\mathrm{QK}}^{\mathrm{T}}\right)}{\sqrt{N/r}}\mathrm{V}\end{array}$

其中，$W_i$表示用于通道投影的线性层，即$W_1\in {\mathbb{R}}^{2C\times LC}$和 { W 2 , W 3 , W 4 } ∈ R C × C \left\{W_{2}, W_{3}, W_{4}\right\} \in \mathbb{R}^{C \times C} {W2​,W3​,W4​}∈RC×C。同时，$R_i$表示实现为带步长的深度可分离卷积（DWConv）的空间缩减器，而 { Q , K } ∈ R G × C G × N r \{\mathrm{Q}, \mathrm{K}\} \in \mathbb{R}^{G \times \frac{C}{G} \times \frac{N}{r}} {Q,K}∈RG×GC​×rN​和$\mathrm{V}\in {\mathbb{R}}^{G\times \frac{C}{G}\times N}$。

SparX引入了两个灵活的超参数，即$\mathbf{S}$和$\mathbf{M}$，这使得网络更容易且更有效地扩展得更深和更宽。此外，DMCA的加入实现了强大的跨层特征交互，从而产生了更强大和鲁棒的表示。

3.2 网络架构

基于所提出的稀疏跨层连接，我们开发了一种名为SparX-Mamba的新型基于Mamba的视觉主干网络，它是一个具有三种不同变体（tiny、small和base）的4阶段层次结构。在第1阶段，由于使用高分辨率输入时计算开销较高，我们不包括任何神经节层。也就是说，第1阶段仅包含两个正常层。在第2阶段，我们将第二层指定为神经节层，该层与其前一层具有层内连接。由于第3阶段有更多层，我们为三个变体设置了不同的超参数。如表1所列，由于tiny模型的深度相对较小，我们将$\mathbf{S}$设置为2，而在small和base模型中，$\mathbf{S}$设置为3，以避免在更深层次的模型中产生过高的计算成本。在第4阶段，我们对tiny和small模型做了特殊处理，将所有层都设置为神经节层，因为此阶段处理的是最低分辨率的特征，在这些模型中产生的计算成本相对较低。对于更复杂的base模型，仅将最后一层设置为神经节层，以防止计算成本过高。同时，为了促进不同阶段之间的通信，每个阶段的第一个神经节层连接到前一阶段最终特征的降采样版本。此外，跨层滑动窗口仅在每个阶段内应用，以防止因特征存储过多而导致内存消耗过高。

4 实验

在本节中，我们从图像分类开始，对代表性视觉任务进行了全面的实验评估。之后，将预训练模型迁移到下游任务，包括目标检测和语义分割。所有实验均在8个NVIDIA H800 GPU上进行。更多实验结果见附录A。

4.1 图像分类

设置。我们使用ImageNet-1K数据集（Deng等，2009），并遵循DeiT（Touvron等，2021）中描述的实验设置，以确保公平比较。我们的方法与许多代表性视觉模型进行了比较：基于CNN的模型（ConvNeXt（Liu等，2022），InceptionNeXt（Yu等，2024b），SLaK（Liu等，2023），UniRepLKNet（Ding等，2024），和MambaOut（Yu和Wang，2024）），基于Transformer的模型（Focal-Transformer（Yang等，2021），PVTv2（Wang等，2022），Swin（Liu等，2021），CSWin（Dong等，2022），和UniFormer（Li等，2023）），以及基于Mamba的模型（PlainMamba（Yang等，2024），ViM2（Behrouz, Santacatterina, 和Zabih，2024），VMamba（Liu等，2024），EfficientVMamba（Pei, Huang, 和Xu，2024），LocalVMamba（Huang等，2024），和MSVMamba（Shi, Dong, 和Xu，2024））。

结果。表2展示了我们的模型在CNN、Transformer和Mamba基方法上的显著优势。例如，SparX-Mamba-T在top-1准确率上分别比ConvNeXt-T和Swin-T高出$1.4\%$和$2.2\%$。与最近的VMamba-T相比，我们的模型在参数更少的情况下实现了$1.0\%$的top-1准确率提升，并且浮点运算数（FLOPs）仅略有增加。此外，即使模型更大，SparX-Mamba仍保持明显优势。同时，我们还在附录A中提供了与更先进视觉模型的更全面比较。总体而言，我们的方法在ImageNet-1K分类任务中表现出竞争性能，并在下游任务中展现出明显优势。

4.2 目标检测和实例分割

设置。为了评估我们的网络架构在目标检测和实例分割方面的性能，我们在COCO 2017数据集（Lin等，2014）上进行了实验。我们使用Mask R-CNN框架（He等，2017），并采用与Swin（Liu等，2021）相同的实验设置。主干网络首先在ImageNet-1K上进行预训练，然后分别在12个周期（$1\times$计划）和36个周期（$3\times$计划+多尺度训练）上进行微调。

结果。如表3所示，我们的模型在目标检测和实例分割方面的性能均优于其他模型。具体来说，在使用$1\times$计划进行微调时，SparX-Mamba-T的$A{P}^b$/$A{P}^m$比VMamba-T分别高出$0.7\%$/$0.4\%$；在使用$3\times$计划进行微调时，SparX-Mamba-T的$A{P}^b$/$A{P}^m$比VMamba-T分别高出$1.3\%$/$1.0\%$。值得注意的是，在使用$3\times$计划时，SparX-Mamba-T的性能甚至超过了具有更多参数和FLOPs的VMamba-S。同时，与其他优秀的视觉主干网络相比，我们的方法也表现出优越的性能。

4.3 语义分割

设置。我们在ADE20K数据集（Zhou等，2017）上进行了语义分割实验。我们采用了两种分割框架：Semantic FPN（SFPN）（Kirillov等，2019）和UperNet（Xiao等，2018）。为确保公平比较，我们使用在ImageNet-1K上预训练的权重来初始化所有主干网络。此外，我们严格遵循先前工作（Liu等，2021；Li等，2023）中概述的相同训练设置。

结果。如表4所示，当将SparX-Mamba集成到S-FPN和UperNet中时，它表现出显著优势。例如，当使用UperNet时，SparX-Mamba的mIoU比VMamba-T高出$1.7\%$/$2.2\%$。small/base模型的性能提升也很明显。同时，当使用S-FPN时，SparX-Mamba-T的mIoU比VMamba-T提高了$2.3\%$，并且比CSWin-T高出$1.3\%$。这些优势在small和base模型中也很明显，有效证明了我们的模型在密集预测任务中的鲁棒性。

4.4 消融实验

设置。为了评估SparX中各个组件的有效性，我们对图像分类和语义分割进行了全面的消融实验。我们首先使用与第4.1节相同的训练设置在ImageNet-1K数据集上训练每个模型变体。然后，我们使用SFPN框架在ADE20K数据集上对预训练模型进行微调，并保持与第4.3节描述的相同训练配置。我们还提供了更多分析实验，详见附录A。

与DenseNet类似网络的比较。基于SparX-Mamba-T，我们进行了以下实验：（1）为了验证我们的稀疏跨层连接，我们在SparX中保持步长$\mathbf{S}=2$，并通过移除滑动窗口并将每个神经节层与所有先前的神经节层和正常层相连来设置跨层连接。这种设计称为“密集神经节连接（DGC-Mamba-T）”。（2）为了验证我们的模型相对于DenseNet（Huang等，2017）的优越性，我们将$\mathbf{S}$设置为1，并移除跨层滑动窗口，从而设计了一个与DenseNet类似的网络。该模型称为“DenseNet风格网络（DSN-Mamba-T）”。

如表5所示，我们观察到与VMamba相比，两种密集跨层连接（DGC和DSN）都产生了明显的计算开销，尤其是在吞吐量显著降低和内存使用量增加方面，这可能会阻碍扩展到更大的模型。相比之下，我们的SparX实现了更显著的改进，同时仅略微降低了吞吐量并稍微增加了内存使用量，证明了其在基于Mamba的模型中的高效性和有效性。

与不同交互方法的比较。基于SparX-Mamba-T，我们将动态多层跨注意力（DMCA）替换为其他多层特征交互的融合方法：（1）移除整个DMCA模块，直接将当前层和选定先前层的特征进行拼接，然后使用一个卷积前馈神经网络（Conv-FFN）层（Wang等人，2022）进行特征融合。使用Conv-FFN进行特征融合保持了与基线模型相当的复杂度。这种设计称为“Concat”。（2）移除DMCA中的跨层分组通道注意力（CGCA）操作。这意味着无需执行空间缩减并计算空间缩减后的查询和键之间的交叉注意力矩阵，但仍需要计算值。然后，将查询和值直接拼接并使用Conv-FFN层进行融合。这种设计称为“w/o CGCA”。（3）将DMCA中的注意力计算替换为空间缩减注意力（SRA）（Wang等人，2021，2022）和跨尺度注意力（CSA）（Shang等人，2023），以比较通道混合和空间混合的性能。（4）将我们的DMCA替换为FcaFormer（Zhang, Hu, and Wang 2023）中提出的空间交叉注意力（SCA），以比较不同多层特征聚合机制的性能。这个版本称为“SCA”。

如表6所示，由于从先前层中动态检索特征，我们的DMCA在跨层特征聚合和选择方面表现出色。具体来说，与我们的DMCA相比，“Concat”和“w/o CGCA”都导致了明显的性能下降。同时，DMCA的性能优于SRA和SCA，可能是因为DMCA提供了动态通道混合，这是对使用SSM执行的空间混合的补充。值得注意的是，一些先前的工作已经证明，结合通道和空间动态可以提高性能（Fu等人，2019；Ding等人，2022a）。相反，SRA和SCA都作为类似于SSM的空间混合器，损害了模型有效表示多维特征的能力。更重要的是，与简单的特征拼接相比，由于二次复杂度，SRA和SCA都产生了更高的内存成本和更低的吞吐量，而我们的DMCA仅略微增加了计算开销，表明其效率更高。

SparX中的超参数设置。我们对所提出的稀疏跨层连接机制中的超参数设置进行了深入探索。以小型模型为基线（即 $\mathbf{S}=2$，$\mathbf{M}=3$），我们构思了以下模型变体：（1）将所有节细胞层替换为普通层，即网络不再有任何跨层连接。这种设计称为“PlainNet”。为了公平比较，我们通过将PlainNet的第3阶段和第4阶段的通道数分别调整为384和576，使本比较中考虑的模型具有相当的复杂度。（2）保持 $\mathbf{M}=3$ 不变，我们将控制节细胞层稀疏性的步长超参数（$\mathbf{S}$）分别设置为1、3和4，以评估步长对最终性能的影响。（3）保持 $\mathbf{S}=2$ 不变，我们通过将 $\mathbf{M}$ 分别设置为2和4，研究了跨层滑动窗口大小的影响。

结果。如表7所示，尽管模型复杂度增加，但具有密集节细胞层（即 $\mathbf{S}=1$）并未导致明显的性能提升。这表明我们稀疏放置节细胞层已经促进了跨层通信。然而，随着 $\mathbf{S}$ 的进一步增加，模型性能逐渐下降，表明节细胞层之间的较大间隔对性能产生了不利影响。请注意，尽管 $\mathbf{S}=3$ 是次优选择，但考虑到更深模型的复杂度增加，我们在小型和基础版本中将 $\mathbf{S}$ 设置为3以提高计算效率。此外，实验表明，在跨层滑动窗口中，先前节细胞层的最佳数量为3，因为将 $\mathbf{M}$ 进一步增加到4并不会带来任何明显的性能提升。

DMCA和DPE的消融实验。以小型模型为基线，我们通过评估以下替代模型设计的性能，进一步研究了DMCA和DPE（Chu等人，2022）的影响。（1）移除DPE模块，所得模型称为“w/o DPE”。（2）通过仅使用GCCA的输出而不将其与 $\mathbf{X}$ 和 ${\mathbf{Y}}_{\mathbf{v}}$ 拼接，研究将GCCA（Z）的输出与查询和值拼接的必要性。这种设计称为“w/o Skip”。（3）消除空间缩减（SR）操作，并在计算GCCA时使用查询和键的原始空间分辨率。这种设计称为“w/o SR”。根据表8中的结果，使用DPE使top-1准确率提高了0.2%，平均交并比（mIoU）提高了0.5%。然后，我们可以发现“w/o Skip”也导致了性能下降，反映了捷径连接的重要性。此外，尽管空间缩减（SR）压缩了空间标记，但它实际上略微提高了性能。一个可能的解释是，在估计图像级特征相似性时，利用压缩信息是足够的。

4.5 通用性分析

设置。为了证明我们跨层连接模式的通用性，我们将SparX融入到具有不同标记混合器架构的网络中进行了实验，这些架构包括自注意力、Mamba中的普通SSM（Single Scale Module，单尺度模块）和ViM中的Bi-SSM（Bidirectional Single Scale Module，双向单尺度模块）。为确保与其他层次化网络进行公平比较，我们遵循了传统设计，即在标记混合器后接一个前馈神经网络（Feed-Forward Network，FFN），标记混合器具有上述多种选择。具体来说，以SparX-Mamba-T作为基本网络，我们将SS2D块替换为Swin Transformer（Liu等人，2021）中提出的移位窗口注意力，以获得一个标记为SparX-Swin-T的网络（即移位窗口注意力 $\to$ FFN），并将第4阶段的通道数从512增加到576，以使其计算成本与Swin-T保持一致。接下来，我们通过将深度增加到$[2,2,27,3]$和通道数增加到$[120,240,396,616]$来扩展SparX-Swin-T的模型规模，以获得一个更大的模型，标记为SparX-Swin-B，其模型复杂度与Swin-B相当。此外，我们还用原始Mamba（Gu和Dao，2024）中的普通SSM模块（即SSM $\to$ FFN）和ViM（Zhu等人，2024）中的Bi-SSM模块（即Bi-SSM $\to$ FFN）替换了VSS中的SS2D模块。为确保公平比较，我们严格保持了其他微块（如移除DPE和使用相同的补丁嵌入层）的相似性。

结果。如表9所示，SparX-Swin-T/B模型在图像分类和密集预测任务上均优于其普通版本，同时保持了相似的计算成本。这表明我们的方法具有通用性，同样适用于基于Transformer的模型。同时，如表10所示，尽管普通SSM和Bi-SSM没有使用VMamba中SS2D所采用的深度可分离卷积（Depthwise Separable Convolution，DWConv），但我们的SparX使普通SSM和Bi-SSM在准确性和计算效率方面均优于VMamba-T。在语义分割任务上，SparX-Mamba-T（SSM）的性能略逊于VMamba-T。这是因为普通SSM中的因果建模未能捕获足够的上下文信息。当使用ViM中的Bi-SSM模块时，平均交并比（mean Intersection over Union，mIoU）显著提高了$0.4\%$，表现出优于VMamba-T的性能。

5 结论

在本文中，我们受人类视觉系统中视网膜神经节细胞（Retinal Ganglion Cell，RGC）层的启发，提出了一种新的跳跃连接策略，命名为SparX。SparX旨在创建稀疏的跨层连接，以增强信息流，并在视觉主干网络中促进特征蒸馏和重用。此外，我们还提出了一种动态多层通道聚合器（Dynamic Multi-layer Channel Aggregator，DMCA），它促进了跨层的动态特征聚合和交互。基于SparX，我们进一步提出了SparX-Mamba和SparX-Swin，两者在一系列具有挑战性的视觉任务中均表现出卓越的性能。

A 附录

A.1 目标检测和实例分割中的更详细比较

我们通过引入基于不同交并比（Intersection over Union，IoU）阈值的指标，提供了目标检测和实例分割方面更全面的性能比较，作为对表3的补充。结果如表11所示，显然我们的模型在所有指标上均优于其他方法。

A.2 与更先进模型的比较

为了呈现更全面的性能分析，我们在图像分类、语义分割和目标检测方面，将我们的SparX-Mamba与高度优化的视觉主干进行了比较。比较的方法包括Conv2Former（Hou等人，2024）、InternImage（Wang等人，2023b）、PeLK（Chen等人，2024）、CMT（Guo等人，2022）、MaxViT-T（Tu等人，2022）、BiFormer（Zhu等人，2023）、MPViT（Lee等人，2022）、NAT（Hassani等人，2023）和CrossFormer++（Wang等人，2023a）。在图像分类任务（表12）上，我们的模型性能比其他基于Mamba的模型更接近最先进的（State-Of-The-Art，SOTA）ViT、CNN及其混合版本。例如，SparX-Mamba-T在Top-1准确率方面与InternImage-T（Wang等人，2023b）相当，并且仅比MaxViT-T（Tu等人，2022）落后$0.2\%$。实际上，我们的模型在准确率上超过了大多数先进的CNN和ViT架构，即SparX-Mamba-T在准确率上比CrossFormer++S（Wang等人，2023a）和NAT-T（Hassani等人，2023）高出$0.3\%$。当模型扩展到更大规模时，也可以观察到类似的现象。更重要的是，尽管SparX-Mamba在图像分类上可能不是表现最好的，但在具有更大输入分辨率的密集预测任务上却表现出优越的性能。如表13和表14所示，虽然SparX-Mamba-T和InternImage-T在图像分类上的性能相当，但SparX-Mamba-T在目标检测和语义分割上均大幅优于InternImage-T。尽管一些高度优化的模型在图像分类上的性能优于SparX-Mamba，但SparX-Mamba在密集预测任务上的表现更佳。例如，CrossFormer+±L在图像分类上优于SparX-Mamba-B，但在语义分割上，后者明显优于前者。请注意，其他基于Mamba的模型并未充分体现出这一优势。SparX-Mamba在密集预测任务上的卓越性能，如附录A.3中可视化所示，归因于其通常更有效的感受野（Effective Receptive Fields，ERFs）。此外，近期研究（Yu和Wang，2024）表明，尽管SSM在图像分类任务上没有性能优势，但在高分辨率输入下的密集预测任务中表现更佳。

A. 3 分析实验

分辨率增加的影响：继VMamba之后，我们进一步评估了模型在不同输入分辨率下的泛化能力。具体来说，我们利用在ImageNet-1K上预训练且输入尺寸为$224\times 224$的模型，对一系列图像分辨率（从$384\times 384$到$1024\times 1024$）进行推理。如表15所示，随着输入分辨率的增加，我们的模型表现出最稳定的性能，显著超越了其他竞争模型。当使用$384\times 384$分辨率时，所有其他模型的性能均有所下降，而我们的SparX-Mamba则显示出明显的性能提升，准确率从$83.5\%$提高到$84.0\%$。此外，随着分辨率的进一步提高，SparX-Mamba的性能下降幅度最小。另外，可以观察到，随着分辨率的增加，VMamba相对于CNN和Transformer模型的优势逐渐减弱，而我们的模型始终保持着显著的性能优势。例如，在$1024\times 1024$输入分辨率下，与Swin相比，VMamba仅实现了$0.4\%$的微小提升，而我们的模型则比Swin高出近$10\%$。这些结果表明，我们的方法能更好地处理包含大量输入标记的情况。

代表性视觉主干网络的速度比较：我们已对代表性视觉主干网络的吞吐量（$T$）进行了全面比较。如表16所示，近期基于Mamba的模型在速度上并未展现出相对于经典的ConvNeXt和Swin模型的优势。这可能是因为深度学习框架尚未得到充分优化，无法像高度优化的GPU矩阵乘法那样在GPU上高效执行SSM（空间移位模块）计算。同时，与卷积和自注意力相比，SSM的并行化能力较低也是另一个影响因素。尽管存在这些原因，但与一些先进的视觉Transformer相比，我们的网络架构在速度和准确性之间仍实现了良好的权衡。例如，SparX-Mamba-T在准确率上比MPViT-S高出显著的0.5%的Top-1准确率，而速度仅下降了10%。与Focal-Transformer相比，我们的模型在速度和性能方面也表现出显著优势。

另一方面，与基于Mamba的模型相比，我们的SparX-Mamba也展现出了更好的速度-准确性权衡。具体来说，从VMamba-T过渡到我们的SparX-Mamba-T时，Top-1准确率显著提高了1%，而吞吐量仅下降了243 imgs/s（从1613下降到1370），同时模型参数也有所减少。相比之下，当VMamba-T演变为VMamba-S时，Top-1准确率提高了1.1%，但速度大幅下降了592 imgs/s（从1613 imgs/s下降到1021 imgs/s）。因此，在Top-1准确率提升约1%的相似情况下，SparX-Mamba-T在吞吐量上相较于VMamba-S（1370 vs. 1021）具有显著优势，且模型复杂度仅约为后者的一半。这清楚地表明，我们提出的架构可以提供更好的速度-准确性权衡。同时，这一优势也适用于更大的模型。具体来说，我们的SparX-Mamba-S在速度、模型复杂度和性能方面都优于VMamba-B，具有更快的速度、更低的模型复杂度和更好的性能。因此，综合考虑速度和准确性，结果清楚地表明，我们的SparX-Mamba在速度-准确性权衡方面相较于VMamba具有显著优势。此外，值得注意的是，EfficientVMamba的第3和第4阶段是完全的卷积层，没有任何Mamba层，因此直接比较计算效率是不公平的。总体而言，我们的SparX-Mamba在计算效率和性能之间取得了出色的平衡，在基于Transformer和Mamba的模型中脱颖而出。

基于Mamba的模型在不同分辨率输入下的速度比较：我们对基于Mamba的最新层次模型在不同分辨率输入图像上的速度和GPU内存消耗进行了全面比较。如表17所示，与其他基于Mamba的模型相比，我们的模型在速度和性能之间实现了最佳权衡。值得一提的是，尽管我们的模型层间连接更多，但与VMamba相比，在GPU内存消耗和延迟方面并未显著增加。这是因为SparX是一种在基于Mamba的模型中创建跨层连接的高效机制。此外，一个有趣的现象是，当在一些更高分辨率上直接进行推理时，我们的小型模型相较于基准模型略有改进。潜在原因可能是模型是在$224\times 224$分辨率下训练的，这限制了其对更大分辨率图像中包含知识的理解能力。结果表明，我们的小型模型在训练过程中学习了更多的分布外知识，从而在性能上略优于基准模型。然而，当模型直接在高分辨率图像上训练时，基准模型在涉及更大分辨率输入的密集预测任务中表现出明显优势。

有效感受野（ERF）分析：为了直观展示我们SparX-Mamba的表示能力，我们可视化了我们的模型以及其他代表性模型（包括ConvNeXt、Swin和VMamba）的有效感受野（Luo等人，2016）。可视化是使用从ImageNet-1K中随机抽取的200张$224\times 224$分辨率的图像生成的。如图4所示，显然，我们的模型在网络的所有阶段都实现了最大的有效感受野。特别是，SparX-Mamba最终阶段的有效感受野几乎涵盖了整个输入图像，展示了其卓越的全局上下文建模能力。

中心核对齐（CKA）分析：由于具有更密集的跨层连接（即表5中的DGC和DSN模型）的模型尽管连接更多且计算复杂度更高，但并未带来性能提升，我们希望通过分析逐层特征和计算层间学习特征模式的相似性，使用中心核对齐（CKA）（Kornblith等人，2019）来发现潜在原因。我们选择了四个代表性模型：VMamba-T、表5中的DGC和DSN模型以及SparX-Mamba-T。在此分析中，使用了每层SSM输出的特征。如图5所示，VMamba-T内每一层学习的模式与附近层非常相似，这一现象也在经典的ViT和CNN模型中被观察到（Raghu等人，2021）。显然，DGC和DSN模型的层可以学习更多独特的模式，这有助于开发更具判别性的表示。然而，当我们将跨层连接稀疏化为SparX-Mamba-T时，与其他方法相比，模型在层间学习了更多样化的特征。因此，模型性能得到了提升。根据这一分析，我们的仿生SparX通过鼓励模型学习更多样化的表示来促进特征交互和蒸馏。尽管DGC和DSN模型促进了特征多样性，但我们的SparX模型进一步增强了这种多样性，这是其实现更好性能的原因。

A. 4 局限性

一些高度优化的卷积神经网络（CNNs）和视觉变换器（ViTs）在ImageNet-1K数据集上的表现优于我们的方法，因为基于Mamba的模型仍处于探索的早期阶段。然而，值得注意的是，我们的SparseCCMamba模型已经实现了接近许多先进CNNs和ViTs的性能，并在附录A.2中提供了更为全面的比较。另一方面，与卷积和自注意力相比，稀疏空间模型（SSMs）具有更复杂的架构。现有的深度学习框架仅支持对SSMs进行有限的优化，导致基于Mamba的模型在速度上不占优势。更多细节见附录A.3。最后，除了本文提供的我们模型的三个版本外，未来我们将探索使用基于Mamba的架构构建大型视觉基础模型的潜力。