VSSD 在图像分类、检测与分割中的应用， 刷新基于 SSM 的模型 SOTA 榜 ！

视觉 Transformer 在计算机视觉领域取得了显著进展，提供了强大的建模能力和全局感受野。然而，它们对计算资源的高需求限制了在处理长序列时的应用性。 为了解决这个问题，状态空间模型（SSM）因其在视觉任务中提供的线性计算复杂性而受到重视。最近，状态空间对偶性（SSD），作为SSM的一种改进变体，在Mamba2中引入以增强模型的性能和效率。 然而，SSD/SSM的固有因果性质限制了它们在非因果视觉任务中的应用。为了克服这一限制，作者提出了视觉状态空间对偶性（VSSD）模型，它采用了SSD的非因果格式。 具体来说，作者建议丢弃隐藏状态与标记之间交互的幅度，同时保留它们的相对权重，这减轻了标记贡献对先前标记的依赖性。结合多扫描策略的运用，作者证明了扫描结果可以整合以实现非因果性，这不仅提高了SSD在视觉任务中的性能，还提升了其效率。 作者在包括图像分类、检测和分割在内的多个基准上进行了大量实验，VSSD超过了现有的基于SSM的最先进模型。 代码和权重可在https://github.com/YuHengsss/VSSD获取。

1 Introduction

近年来，由Vision Transformer（ViT）[9]引领的视觉 Transformer （vision transformers）在计算机视觉领域取得了巨大成功。得益于注意力机制的全局感受野和强大的信息建模能力，基于视觉 Transformer 的模型在分类[7]、检测[32]和分割[66]等各项任务中均取得了显著进展，超越了经典的基于CNN的模型。然而，注意力机制的二次计算复杂度使得它在处理长序列的任务中资源消耗巨大，这限制了其更广泛的应用。

近期，状态空间模型（State Space Models, SSMs），特别是以Mamba[13]为例，引起了研究者的广泛关注。特别是S6模块，它具有全局感受野，并且与序列长度成线性关系，提供了一个有效的替代方案。先驱的视觉Mamba模型，如Vim[68]和V Mamba[34]，将SSMs应用于视觉任务。之后，许多变体被提出，这些变体通过不同的扫描路径将2D特征图展平为1D序列，使用S6模块进行建模，然后在多个扫描路径中整合结果。这些多扫描方法提高了SSMs在视觉任务中的性能，取得了与基于CNN和ViT方法相媲美的成果。最近，Mamba2[6]对S6模块进行了进一步的改进，提出了状态空间对偶性（State Space Duality, SSD）的概念。Mamba2将状态空间转移矩阵\mathbf{A}视为标量，扩展了状态空间维度，从而提高了模型性能以及训练和推理的效率。然而，在将SSD/SSMs应用于视觉任务中存在一个主要问题，即图像数据本质上是非因果的，而SSD/SSMs具有固有的因果属性。另一个问题是，将2D特征图展平为1D序列破坏了各区块之间固有的结构关系。作者在图1（a）中提供了一个说明，以更直观地理解这两个问题。在这个例子中，展平的1D序列中的中心标记仅限于访问之前的标记，无法整合后续标记的信息。此外，在2D空间中与中心标记相邻的标记1在1D序列中被远远地放置，破坏了自然的结构关系。先前解决方案[34, 28]中的一种常见做法是增加对非因果视觉特征的扫描路径，这在一定程度上缓解了这两个问题。鉴于这些观察，一个重要的问题产生了：与多扫描方法相比，是否存在一种更有效且高效的方法将SSD应用于非因果视觉数据？

为了回答这个问题，作者对SSD的分析表明，将矩阵\mathbf{A}视为标量，可以简便地将SSD转换为非因果性和位置无关的方式，作者将其称为非因果SSD（NC-SSD）。具体来说，作者不是使用\mathbf{A}来确定要保留的隐藏状态的比例，而是用它来指示当前标记对隐藏状态的贡献程度。在这种情况下，每个标记的贡献变得自指。基于这一特性，作者证明了SSD中的因果 Mask 可以自然地移除，无需特定的扫描路径。这一观察促使作者开发了一种非因果格式的SSD，其中可以导出一个全局隐藏状态以替代先前的按标记隐藏状态，不仅提高了准确性，还增强了训练和推理的速度。与先前主要缓解SSM因果限制的多扫描方法[68, 34]不同，作者提出的NC-SSD还解决了将2D特征图展平为1D序列会破坏相邻标记连续性的问题。除了NC-SSD，还包括与其他技术如标准自注意力机制结合以及重叠下采样等探索。基于这些技术，作者提出了视觉状态空间二分法（VSSD）模型，并在图1（b）和（c）中展示了与基于CNN、ViT和SSM的方法相比，其在有效性和效率上的优势。具体而言，与最近提出的基于SSM的VMamba[34]相比，在ImageNet-1K数据集[7]上，作者的VSSD模型的top-1准确率提高了大约1%，同时保持了相似的计算成本。此外，作者的模型在准确性与延迟曲线上也始终领先。除了在性能与效率之间取得更好的平衡之外，VSSD的另一大亮点是训练速度。例如，与普通的SSD或多扫描SSD（例如，双向扫描的Bi-SSD）相比，作者提出的模型分别将训练速度提高了近20%和50%。

总之，作者的贡献有两方面。首先，作者分析了状态空间的对偶性，并证明它可以无缝地转换到非因果模式。基于这一洞见，作者引入了NC-SSD，它保留了原始SSD的全局感受野和线性复杂度优势，同时融入了固有的非因果属性，实现了训练和推理效率的提升。其次，以NC-SSD为基础组件，作者提出了VSSD模型，并通过广泛的实验验证了其有效性。在相似的参数和计算成本下，作者的VSSD模型在分类、目标检测和分割等多个广泛认可的基准测试中，超越了其他基于SSM的现有最优（SOTA）模型。

2 Related Work

视觉 Transformer (Vision Transformers) 视觉 Transformer （ViTs）的引入为先前由卷积神经网络（CNNs）主导的计算机视觉领域注入了新的活力。然而，ViTs中的自注意力机制的二次计算复杂度在处理高分辨率图像时带来了重大挑战，需要大量的计算资源。为了解决这个问题，提出了不同的解决方案，包括分层架构、窗口化注意力[36; 21; 51; 67]以及自注意力的变体[53; 57; 63]。同时，线性注意力[29; 4; 40; 19]通过改变自注意力中 Query 、键和值的计算顺序，成功地将计算复杂度降低到线性规模。尽管有了这些进步，线性注意力的性能仍然低于二次自注意力[52]及其变体[21; 11; 67]。

状态空间模型 (State Space Models) 状态空间模型（SSMs）由于其全局感受野和线性的计算复杂度，越来越多地吸引了研究者的关注。Mamba [13]，作为SSMs的一个突出例子，引入了S6块，在自然语言处理（NLP）基准测试中取得了与 Transformer 相当或更好的性能。后续的努力[39; 28; 10; 60; 2; 31; 62; 43]探索了将S6块适配到视觉任务，与CNN和基于ViT的模型相比，取得了具有竞争力的结果。在开发基于Mamba的视觉模型时，一个核心挑战是将Mamba块的固有的因果性质适配到非因果的图像数据。最直接的方法包括使用不同的扫描路径将2D特征图展平为1D序列，然后使用S6块进行建模和整合。受到这些考虑的启发，各种扫描路径已被采用并被证明是有效的，多项研究[68; 34; 28; 39; 46]提供了证据。最近，Mamba2 [6]强调了状态空间模型与结构化 Mask 注意力之间的重大重叠，认定它们是彼此的对偶，并引入了状态空间对偶性（SSD）的概念。在此基础上，作者证明，通过简单的变换，可以将SSD转换为非因果模式，无需特定的扫描路径。

3 Method

Preliminaries

状态空间模型。 经典的状态空间模型（SSM）用于描述连续系统的动态，将输入序列 转换为潜在空间表示，然后利用这种表示生成输出序列 。

Non-Causal State Space Duality

为了使SSM适应图像数据，首先需要将2D特征图展平成一个1D的 Token 序列，然后对这些 Token 进行顺序处理。由于SSM的因果性质，每个 Token 只能访问之前的 Token ，因此信息传播本质上是单向的。这种因果属性在处理非因果图像数据时会导致性能不佳，这一发现得到了先前研究的证实。

此外，将2D特征图展平成1D序列会破坏其固有的结构信息。例如，在2D图中相邻的 Token 可能会在1D序列中相隔甚远，导致在视觉任务上的性能损失[17]。由于SSD是SSM的一种变体，将SSD应用于视觉任务时，会面临与SSM相似的挑战：

• 挑战1：模型的因果性质限制了信息流，阻止了后续 Token 对先前 Token 的影响。
• 挑战2：将2D特征图展平成1D序列，在处理过程中破坏了 Patch 间的固有结构关系。

Vision State Space Duality Model

块设计。 为了在Mamba2中增强适用于视觉应用的SSD块，作者在将SSD简单替换为NC-SSD以开发作者的视觉状态空间二元性（VSSD）块的基础上，实施了几项修改。在构建NC-SSD块时，将因果卷积1D替换为深度卷积（DWConv），其核大小为三，与先前的视觉Mamba工作[34；28]保持一致。此外，在NC-SSD块之后集成了一个前馈网络（FFN），以促进通道间的信息交换，并与经典视觉 Transformer [9；36；50]的既定实践保持一致。此外，在NC-SSD块和FFN之前加入了一个局部感知单元（LPU）[18]，增强了模型对局部特征感知的能力。不同块之间也实现了跳跃连接[23]。VSSD块的架构在图4的下半部分展示。

与自注意力结合。 Mamba2表明，将SSD与标准的多头自注意力（MSA）结合可以带来额外的改进。同样，作者的模型也融入了自注意力。但是，与在网络的各个部分均匀穿插自注意力的Mamba2不同，作者只在最后一个阶段将NC-SSD块专门替换为自注意力模块。这一修改利用了自注意力在处理高级特征方面的强大能力，正如先前在视觉任务中的工作[33, 42, 11]所证明的。

重叠下采样层。 由于分层视觉 Transformer [36]和视觉状态空间模型[34]主要使用非重叠卷积进行下采样，最近的研究[21, 56]表明，重叠下采样卷积可以引入有益的归纳偏置。因此，作者采用了MLLA[19]中的重叠卷积方法。为了保持参数数量和计算FLOPs的可比性，作者相应地调整了模型的深度。

总体架构。 作者按照上述方法开发VSSD模型，其架构如图4所示。借鉴先前工作[36, 37, 34]中建立的视觉 Backbone 网络的设计原则，作者的VSSD模型被构建为四个层次化阶段。前三个阶段采用VSSD块，而最后一个阶段则整合了MSA块。VSSD变体的详细架构在表1中展示。

4 Experiment

4 Experiment

在本节中，作者将详细描述实验的设计、实施和结果分析。作者将首先介绍实验设置，包括数据集的选择、评估指标以及所采用的方法。随后，作者将讨论实验过程中观察到的关键现象，并对实验结果进行详细分析。最后，作者将探讨这些结果对相关领域的潜在影响和意义。

Classification

配置。作者的实验使用ImageNet-1K数据集[7]，与之前的研究[36, 34]的方法论一致。每个模型都要经历一个跨越性能评估。表2展示了在ImageNet-1K数据集[7]上，作者的VSSD模型与CNNs、ViTs和其他基于SSM的框架的对比情况。拥有14M参数和2.3G FLOPs的VSSD-M模型，其top-1准确度为82.5%，超过了价格相近的NAT-M[21] 0.7%。在小型和中小型模型类别中，VSSD一致性地超越了其竞争对手。特别是，拥有24M参数和4.5G FLOPs的VSSD-T模型，其准确度为83.7%，比VMambaV9-T[34]高出1.2%。对于小型模型变体，拥有40M参数和7.4G FLOPs的VSSD-S，其准确度为84.1%，比LocalVMamba-S[28]高出0.4%。在基础尺寸的变体中，拥有89M参数和16.1G FLOPs的作者的VSSD-B，其准确度为84.7%，比VMambaV9-B高出0.8%。

除了定量比较之外，作者还对训练前后各种模型的有效感受野（ERF）进行了比较分析，包括基于CNN的ResNet50[23]和ConvNeXt-Tiny[37]，基于注意力的Swin-Tiny[36]和DeiT-Small[50]，以及基于SSM的VMamba-Tiny[34]，以及作者的VSSD-Tiny。使用[38]提出的方法绘制了ImageNet-1K验证集中随机选择的50张分辨率为1024x1024的图像中央像素的ERF。为了证明所提出的NC-SSD的有效性，在第3.3节中讨论的混合自注意力技术和重叠下采样层等技术在作者的VSSD模型此分析中未使用。值得注意的是，只有作者的VSSD和DeiT在训练前后都表现出全局感受野。在训练后，VMamba出现了明显的十字形衰减，而作者的方法有效地消除了token间隔对信息贡献的影响。

Object Detection and Instance Segmentation

配置。 对VSSD模型的评估在Mask R-CNN框架 [22] 下使用MS COCO数据集 [32] 进行目标检测和实例分割任务。所有实验均使用MMDetection库 [1] 进行。与先前研究 [36; 34] 一致，在训练阶段，调整图像以使短边为800像素，而长边不超过1333像素。优化采用AdamW优化器，学习率为0.0001，批量大小为16。"1x"标准训练计划下，在第8和第11个周期将学习率降低10倍；而在扩展的"3x + MS"计划中，在第27和第33个周期以同样的比例降低学习率。性能评估。 表3详细比较了作者的模型与已确立的CNN、ViT和其他基于SSM的模型的性能。作者的VSSD模型在各种配置中表现出色。值得注意的是，作者的VSSD-T模型显示出显著优势，相较于Swin-T [36]，在边界框AP上高出+4.2，在 Mask AP上高出+3.3。在扩展的"3x"训练计划下，VSSD-T仍然一致地优于众多竞争者。

Semantic Segmentation

配置。 与Swin [36]和VMamba [34]中描述的方法一致，作者的实验采用了UperHead [58]框架，并使用在ImageNet上预训练的 Backbone 网络进行初始化。训练过程包括160K次迭代，批次大小为16，使用MMSegmentation库[5]执行。主要实验采用标准的输入分辨率512\times 512。为了进一步评估作者模型的鲁棒性，实施了多尺度（MS）测试。优化采用AdamW优化器，学习率设定为6\times 10^{-5}。

性能评估。 表4详细展示了作者模型及其竞争对手的性能指标，包括单尺度和多尺度测试场景。特别是，在Tiny模型类别和单尺度测试的背景下，作者的VSSD模型表现出色，分别超过了Swin、ConNeXt和VMamba模型在Tiny变体上的性能，提高了+3.5、+1.9和+0.6 mIoU。### 消融研究

为了验证所提出模块的有效性，作者在VSSD-Micro模型上进行了详细的消融实验。以SSD块作为标记混合器和块状下采样器（例如，在stem中使用4\times 4卷积核和步长为4的卷积），遵循Swin [36]和vallina VMamaba [34]，作者建立了基准配置，具体见表5的第一行。为了吞吐量测试，作者使用了A100-PCIE-40G GPU，批次大小为128，并使用FP16精度。

不同的SSD机制。 在作者对标记混合器的消融研究中，作者探索了SSD的不同扫描路径。具体来说，作者引入了Bi-SSD，将通道分成两部分，并将一部分反转以创建反向扫描序列。这些具有相反扫描路径的序列在SSD块之后进行连接。如表5所示，作者的NC-SSD模型在top-1准确度上分别比原始SSD和Bi-SSD高出0.6%和0.2%。此外，训练和推理的吞吐量也有所提高，与Bi-SSD方法相比，NC-SSD将训练吞吐量提高了近50%。

混合架构与重叠下采样器。 表5的最后两行展示了将标准注意力纳入最后阶段并使用重叠下采样器的有效性。具体来说，将NC-SSD在最后阶段替换为标准注意力，在略微减少参数的同时，准确度提高了0.2%。将切块下采样器替换为重叠卷积方式，在FLOPs增加0.2G的同时，准确度提高了0.7%。为了保持近似的参数量，作者将层配置从[2,4,8,4]调整为[2,2,8,4]。

参数m的影响。 方程10将NC-SSD概念化为线性注意力的一个变体，它引入了一个额外的权重向量\mathbf{m}。图3直观展示了\mathbf{m}如何选择性地强调前景特征。为了定量评估\mathbf{m}的影响，作者在100个训练周期（其中前5个周期用于 Warm up ）内进行了带与不带此组件的NC-SSD块的实验。表6的结果显示，\mathbf{m}对模型性能有显著影响。

在没有\mathbf{m}的情况下，作者的实验表明，模型在训练过程中会出现不稳定现象，甚至导致崩溃。这种不稳定性在大型模型中尤为明显。作者报告了训练崩溃前达到的最高准确率，用\dagger标记。对于微型模型，最佳准确率仅为32.6%。对于小型模型，训练在第一个时期就崩溃了。作者假设这种不稳定性产生的原因是，在没有通常在线性注意力方法[4, 40, 44]中使用的归一化技术的情况下，特征的大小会急剧增加，从而导致崩溃。

5 Limitations

尽管所提出的VSSD模型在ImageNet-1K上的表现优于其他基于SSM的模型，但与这些基于SSM的模型相比，VSSD在下游任务[32, 66]上的性能提升是有限的。当与最先进的视觉 Transformer 变体[11, 57, 45]进行比较时，在下游任务上的性能仍存在较大差距。此外，本文缺乏涉及更大模型和更广泛数据集的实验，例如使用ImageNet-22K基准[7]的实验。因此，所提出的VSSD模型的可扩展性是未来探索的一个重要领域。

6 Conclusion

总之，作者的研究提出了NC-SSD，它通过改变矩阵A的作用并取消因果掩码，对SSD进行了重新定义。这些调整使得模型能够切换到非因果模式，极大地提高了准确性和效率。

大量的实验验证了NC-SSD相较于原始SSD及其多扫描变种的优越性。此外，通过结合混合标准注意力机制和重叠下采样等技术，我们的VSSD模型在多个广泛采用的基准测试中，与成熟的卷积神经网络（CNNs）、视觉变换器（ViTs）和视觉状态空间模型（Vision SSMs）相比，展现出了相当或更优的性能。

参考

[1].VSSD: Vision Mamba with.