点击下方卡片，关注 「AI视界引擎」 公众号
尽管视觉 Transformer （ViTs）在各项视觉任务中取得了显著成就，但它们受到二次复杂度的限制。近来，状态空间模型（SSMs）因其全局感受野和对输入长度的线性复杂度而受到广泛关注，在自然语言处理和计算机视觉等领域展现出巨大潜力。
为了提高SSMs在视觉任务中的性能，人们广泛采用了多扫描策略，但这导致了SSMs的显著冗余。
为了在效率和性能之间取得更好的平衡，作者分析了多扫描策略成功背后的原因，其中长距离依赖发挥着重要作用。基于这一分析，作者引入了多尺度视觉曼巴（MSVMamba），以有限的参数保留SSMs在视觉任务中的优势。它对原始特征图及其降采样特征图采用多尺度2D扫描技术，这不仅有利于长距离依赖学习，还能降低计算成本。
此外，作者还整合了卷积前馈网络（ConvFFN）来解决通道混合不足的问题。
作者的实验表明，MSVMamba具有高度竞争力，其中MSVMamba-Tiny模型在ImageNet上达到82.8%的top-1准确率，在COCO上使用Mask R-CNN框架，1x训练计划，达到46.9%的边界框mAP和42.2%的实例mAP，在ADE20K上进行单尺度测试达到47.6%的mIoU。
1 Introduction
在计算机视觉领域，特征提取在各种任务中起着关键作用，这些任务从图像分类到更复杂的应用，如检测和分割。传统上，卷积神经网络（CNNs）一直是特征提取方法的主干，因其线性缩放复杂性和捕捉局部模式的能力而受到重视。然而，CNNs在封装全局上下文方面往往力不从心，这一局限性在需要全面理解整个视觉场的任务中变得越来越明显。
相比之下，视觉 Transformer （ViTs）作为一种有力的替代品出现，它们固有的全局感受野允许直接捕获图像内的长距离依赖。尽管它们有优势，但ViTs在输入大小方面的二次缩放复杂度限制了它们在下游任务中的应用，如目标检测和分割，在这些任务中效率至关重要。
最近，基于状态空间模型（SSM）的方法[12; 42; 8]因其能够结合全球感受野和线性缩放复杂度的优点而受到关注。特别是，Mamba[10]提出了一种硬件感知和输入依赖的算法，显著提高了SSMs的性能和效率。
受到Mamba成功的启发，大量研究开始利用其在视觉任务中的优势，开创性工作如ViM[61]和VMamba[31]。
Mamba [10]开发的S6模块最初是为NLP任务设计的。为了使S6适应视觉任务，首先将图像划分为多个块，然后沿着扫描路径将块展平成序列。为了适应图像数据的非因果性质，视觉任务广泛采用多扫描策略，例如ViM [61]通过在前向和反向对序列求和来增强序列，以及VMamba [31]整合水平扫描和垂直扫描。然而，与可以拥有数十亿参数的NLP模型不同，当前的视觉 Backbone 网络总是考虑计算成本，即准确性与效率之间的权衡。
这种对模型尺寸的限制内在地制约了视觉任务中SSMs的长距离建模能力。以ViM-Tiny [61]为例，将cls标记置于序列中间比将其置于末端能显著提高结果。
这表明中心位置补偿了模型整合远距离信息能力的限制，突显了参数受限的视觉模型处理长距离依赖关系的困难。作者将这个问题称为长距离遗忘问题。
在这项工作中，作者分析了[61, 31]中的多扫描策略如何帮助减轻这个问题。与单扫描策略相比，多扫描策略允许长距离衰减在2D图像的各个方向上表现。
然而，增加的扫描路径带来了多倍的计算，显著增加了冗余并限制了效率。因此，作者旨在在视觉任务中追求Mamba性能与效率之间更好的权衡。
最直接有效的方法来解决长距离遗忘问题就是缩短序列长度，这可以通过下采样特征图来实现。然而，将所有的扫描路径放置在一个下采样的特征图上可能会导致细粒度特征的丢失和下游任务性能的下降。通过不同扫描的可视化，作者展示了不同的扫描路径其衰减率可能不同，这促使作者开发出多扫描的分层设计。
在这项工作中，作者提出了一种多尺度的2D (MS2D)扫描策略，以有限的计算成本减轻长距离遗忘问题。具体来说，作者将SS2D [31]的扫描方向分为两组：一组保留原始分辨率并由S6块处理，而另一组则进行下采样，由S6块处理后进行上采样，这不仅缩短了序列长度以学习长距离依赖，同时也减轻了冗余。
在具有分层架构的VMamba的基础上，作者在块内再引入另一分层设计，从而在层次内创建了一个层次。此外，作者在每个块内引入了卷积前馈网络（ConvFFN），以增强模型在通道信息交换和局部特征捕获方面的能力。
作者进行了广泛的实验来验证MSVmamba在包括图像分类、目标检测和语义分割在内的一系列任务上的有效性。在ImageNet-1K [3]数据集上的详细对比结果如图1所示。
具体来说，MSVMamba-Tiny相较于VMamba-Tiny取得了显著的提升，在GFLOPs上同时减少了17%的计算需求的同时，将Top-1准确率提高了0.6%。此外，作者的MSVMamba模型在COCO [29]和ADE20K [59]数据集上的目标检测和语义分割任务中的性能也超越了现有的基于SSM的模型。
2 Related work
Generic Vision Backbone
通用视觉 Backbone 网的演变显著塑造了计算机视觉的格局，从CNNs [26, 41, 19, 55, 22, 20, 43, 34] 发展到ViTs 。在深度学习的早期时代，CNNs一直是基于视觉的模型的基础，主导了视觉任务。经典的CNNs如AlexNet [26]，VGG [41]和ResNet [19]为后续众多创新[55, 22, 20, 43, 21, 34, 4, 50]铺平了道路。
这些设计显著提高了网络从视觉数据中捕捉复杂模式和特征的能力，从而广泛提升了视觉任务的性能。受到在自然语言处理中成功的transformers [48]启发，Vision Transformer [6]成为了传统CNNs在视觉相关任务上的有力竞争者。ViT通过将图像分割成块并使用自注意力机制处理这些块，重新构思了图像处理。这种创新方法使模型能够理解整个图像的全局依赖性，这是理解复杂视觉数据的一个重大飞跃。
然而，ViT架构需要相当的计算资源和大规模的 数据集进行有效训练。此外，其性能与输入序列长度紧密相关，表现出二次复杂度，可能会增加处理成本。对此，后续研究致力于开发更有效的训练策略，分层网络结构，改进的空间注意力机制以及基于卷积的设计来解决这些问题。
State Space Models
状态空间模型（SSMs）因其与输入序列长度成线性增长的计算复杂性和固有的全局感知属性而越来越多地受到研究者的关注。为了减少SSMs的计算资源消耗，S4 [12] 引入了对角线结构，并将其与对角线加低秩方法相结合构建结构化SSMs。随后，S5 [42] 和 H3 [8] 通过引入并行扫描技术和提高硬件利用率进一步提升了基于SSM模型的效率。Mamba [10] 接着引入了S6模块，通过融合数据相关参数来修正之前SSM模型的线性时不变（LTI）特性，在大规模数据集上展示了优于 Transformer 的性能。在视觉任务领域，S4ND [36] 首次将SSM模型应用于视觉任务，将视觉数据视为1D、2D和3D信号。U-Mamba [35] 将CNN与SSM结合用于医学图像分割。ViM [61] 和 VMamba [31] 将S6模块整合到视觉 Backbone 设计中，采用多个扫描方向以适应图像数据的非因果性，在与ViTs和CNNs的竞争中取得了有竞争力的结果。受到ViM和VMamba成功的启发，大量基于Mamba的工作在各个视觉任务中涌现，包括视觉 Backbone 设计、医学图像分割和视频分类[27]，展示了基于SSM方法在推进计算机视觉领域的潜力。
3 Method
在本节中，作者首先在第3.1节总结了状态空间模型。
随后，在第3.2节中，作者对现有视觉Mamba模型中的多扫描策略进行了深入分析。
分析之后，第3.3节通过引入多尺度2D（MS2D）扫描策略来解决冗余和长距离依赖问题。
最后，在第3.4节中详细介绍了多尺度状态空间（MS3）块的集成，该块结合了MS2D技术以及一个ConvFFN。
基于MS3块，作者在不同的尺度上开发出了各种模型配置，展示了作者所提出方法的适应性和可扩展性。
Preliminaries
状态空间模型。 经典的状态空间模型（SSM）表示一个连续系统，它将输入序列 映射到一个潜在空间表示 ，然后基于此表示预测输出序列 。从数学上讲，SSM 可以描述如下：
其中 , 和 是可学习的参数。
离散化。 为了在深度学习框架内适配连续状态空间模型（SSM），实施离散化操作至关重要。通过引入时间尺度参数 并采用广泛使用的零阶保持（ZOH）作为离散化规则，可以推导出 和 的离散化版本（分别表示为 和 ），据此，方程式 1 可以改写为离散形式：
其中 表示单位矩阵。之后，方程式 2 的过程可以以全局卷积方式实施：
其中 是卷积核。
选择性状态空间模型。Mamba [10] 引入的选择性状态空间（S6）机制使参数 和 成了输入依赖的，从而提升了基于 SSM 的模型的性能。在使 和 成输入依赖后，方程式 3 中的全局卷积核可以重写为：
Analysis of Multi-Scan Strategy
在处理图像数据时使用S6模块，将2D特征图 展平成一个1D图像标记序列，表示为 。根据方程式4和方程式2，S6中第个标记对第个标记（）构建的贡献可以表示为：< p=""></n$）构建的贡献可以表示为：<>
通常，学到的 是负数，这使模型倾向于优先考虑最近标记的信息。因此，随着序列长度的增加，方程式5中的指数项 快速衰减，导致远距离标记的贡献最小。作者将此称为长距离遗忘问题，这也在最近将S6应用于视觉任务的研究[14]中被观察到。尽管可以通过增加模型参数数量和深度来缓解这个问题，但这种调整引入了额外的计算成本。此外，S6模块的因果属性确保了标记之间的信息只能单向传播，防止早期标记访问后续标记的信息。
图像固有的非因果性质使得将S6模块直接应用于视觉相关任务效果不佳，正如ViM [61]所确定的。为了缓解这一限制，ViM [61] 和 VMamba [31] 引入了通过不同方向扫描图像特征然后整合这些特征的方法论。通常，沿着某条扫描路径更新的标记，表示为 ，其中 指示坐标，可以通过以下方式获得：
在等式6中，代表将2D特征图转换为1D序列的变换，而表示在变换后的1D序列中的对应索引。之后，将来自多扫描路径的结果相加，产生增强特征，可以表示为：
图2：沿水平、垂直扫描路径的衰减及其比率的说明。
不同的扫描路径对应于等式6中的两种不同变换，分别是扁平化和带翻转的扁平化变换。同样，V Mamba [31] 通过结合水平和垂直扫描方向，扩展了基本的双向扫描，产生了四个不同的扫描路径。此外，多扫描策略还通过最小化标记之间的有效距离，减轻了长距离遗忘问题。对于坐标为和的标记，该策略采用多个扫描路径，每个路径都可能改变它们的相对位置。这些路径上的最小距离由给出，其中表示通过第次扫描生成的序列中标记之间的距离。通过减少这个距离，多扫描策略减少了远距离标记之间影响衰减，从而增强了模型维持和利用长距离信息的能力。
为了更直观地展示多扫描策略与长距离衰减之间的关系，作者在V Mamba-Tiny中沿着水平方向和垂直方向可视化了指数项，相对于图1(a)和图1(b)中的中心标记。具体来说，作者从ImageNet [3]验证集中随机选择50张图像，并计算了这些图像和特征维度在最后一层最后一个阶段的扫描路径上的平均衰减。作者使用更高的输入分辨率来增强可视化的质量。
根据这些观察，V Mamba中多扫描策略的成功可以归因于它减轻了图像数据的非因果性质和长距离遗忘问题。然而，随着扫描路径数量的增加，计算成本也线性上升，引入了计算冗余。在图1(c)中，作者说明了图1(a)与图1(b)的最大比率以及反之。而在图1(d)中，作者展示了图1(c)的二值化版本，应用了一个10的阈值，覆盖了整个图形的40%以上。这种现象表明，不同扫描路径之间变化的衰减率导致某些路径主导衰减动态，这也可归因于长距离遗忘问题。主导扫描路径的存在意味着某些扫描对信息保持的贡献显著大于其他扫描，导致多扫描策略中的计算冗余。
Multi-Scale 2D Scanning
正如在上一个小节中讨论的，随着扫描距离的增加， Token 的贡献衰减。缓解长距离遗忘问题的最有效且直接的方法是减少 Token 的数量。同时，由于S6块的计算复杂性与 Token 数量线性相关，减少 Token 数量也提高了效率。因此，解决上述问题的另一种方法是将在下采样的特征图上应用多扫描策略。然而，将所有扫描设置在下采样特征图上将会忽略细粒度特征，并导致不可避免的信息丢失。因此，沿完整分辨率的特征图进行扫描也是必要的。
图3：多尺度2D-选择性扫描在图像上的说明
基于这些考虑，作者提出了一种简单而有效的多尺度2D扫描（MS2D）策略，如图3所示。作者的方法从生成不同尺度的分层特征图开始，通过应用具有不同步长值的深度卷积（DWConv）来实现。这些多尺度特征图随后在V Mamba内的四个不同扫描路线上进行处理。具体来说，作者使用步长为1和s的DWConvs来获得特征图 和 。之后， 和 通过两个S6块处理如下：
其中 是将2D特征图转换为在SS2D中使用的1D序列的转换， 是处理后的序列。这些处理后的序列被转换回2D特征图，并对下采样特征图进行插值以进行合并：
其中 是 的逆变换， 是通过MS2D增强的特征图。
下采样操作将序列长度减少了 倍，这也将等式5中的 Token 距离缩短了 倍，从而缓解了长距离遗忘问题。由于单个S6块的复杂度为 [10]，其中 表示SSM维度，用MS2D替换SS2D将四个扫描的总序列长度从 降低到 ，从而提高了效率。实际上，下采样率设置为2。值得注意的是，来自 的序列由相同的S6块处理。这种方法在有效减少参数数量的同时，保持了使用不同扫描路线的多个S6块相同的准确性。
为了更好地说明长距离遗忘问题的缓解，作者还提供了实证证据，如图4所示。作者比较了VLambda中的SS2D和作者MS2D的扫描路线上的衰减，重点关注与图2相同配置的最后一个 Token 。下采样特征中的衰减图被插值回。正如观察到的，下采样地图上扫描路线的衰减率显著缓解，增强了捕获全局信息的能力。
Overall Model Architecture
在本研究中，作者通过用作者的多尺度状态空间（MS3）块替换VLambda框架的VSS块，扩展了其功能。
MS3块的结构框架如图5所示，包括多尺度视觉空间状态（MSVSS）组件和卷积前馈网络（ConvFFN）。MSVSS组件是通过将VLambda中的视觉状态空间框架适应而来，用MS2D替换SS2D，以在单层中进一步引入层次设计。此外，根据相关文献[24; 38]，在多尺度2D扫描之后集成了挤压激励（SE）[21]块。与之前视觉Mamba架构中关注 Token 混合的研究[61; 31; 24]不同，作者的设计引入了通道混合器以增强不同通道间信息流的增强，与典型视觉 Transformer 的结构范式保持一致。与之前的研究一致，采用由深度卷积和两个全连接层组成的卷积前馈网络（ConvFFN）作为通道混合器。为了确保ConvFFN和MSVSS块在LeViT [9]之后消耗大约相同的FLOPs，作者采用了适度的通道扩展比2。在MS3块中融合MSVSS和ConvFFN时，作者对块的数量进行了细致调整，以确保可比较的计算预算，从而进行公平的比较。
为了实证验证作者提出的修改的有效性，作者在不同的尺度上引入了模型变体，其规格细节见附录6。这些变体，即Nano、Micro和Tiny，分别以其参数计数6.9M、11.9M和33.0M为特征。在计算开支方面，这些模型分别需要0.9、1.5和4.6 GFLOPs，展示了一种可伸缩的模型设计方法，能够适应不同的计算限制。
4 Experimental Validation
ImageNet Classification
设置. 作者的模型在ImageNet-1K数据集[3]上进行训练和测试。与之前的工作[33; 31; 24]保持一致，所有模型都经过300个周期的训练，前20个周期用于预热。训练在8个GPU上使用1024的批量大小。作者采用AdamW优化器，将betas设置为(0.9, 0.999)，动量设置为0.9。学习率通过余弦衰减调度程序进行管理，初始速率为0.001，结合权重衰减0.05。此外，作者利用指数移动平均（EMA）并实施标签平滑，因子为0.1，以提高模型性能和泛化能力。在测试期间，图像以中心裁剪的方式，尺寸为224224。在处理多尺度扫描路径时，作者凭经验选择从左上到右下的路径来处理全分辨率特征图，而其他三个扫描路径负责扫描降采样特征图。
结果. 表1展示了作者的MSV-Mamba模型与已建立的CNNs、ViTs和基于SSM的模型在ImageNet-1K上的对比。具有700万参数和0.9G FLOPs的MSVMamba-N达到了77.3%的top-1准确率，超过了类似成本的基于卷积的RegNetY-800M和基于SSM的EfficientVMamba-T。具有1200万参数和1.5G FLOPs的MSVMamba-M达到了79.8%的准确率，比ViM-T高出3.7%。MSVMamba-T模型在拥有3300万参数和4.6G FLOPs的情况下，达到了82.8%的准确率，比VMamba-T高出0.6%，且计算成本大大降低。这些结果突显了MSVMamba的高效性和可扩展性，为高准确性和资源效率的模型设计提供了一个强大的选择。
Object Detection
设置。 作者在MSCOCO [29] 数据集上使用Mask R-CNN [18] 框架评估作者的MSVMamba，用于目标检测和实例分割任务。遵循之前的工作 [33; 31]，作者使用在ImageNet-1K上预训练的 Backbone 网络进行初始化。为了公平比较，作者采用了标准的1（12个周期）和3（36个周期）训练策略，并采用多尺度（MS）训练。
结果。 表2展示了作者的方法与CNNs、ViTs和基于SSM的模型的性能对比。作者的模型在各种变体和训练设置中一致地超越了其他模型。特别是，MSVMamba-T在1计划下比Swin-T在边界框AP上高出+4.2，在 Mask AP上高出+2.9，并且在3计划下在边界框AP和 Mask AP上也显示出改进。
Semantic Segmentation
设置。 与Swin [33]和VMamba [31]中使用的方法论一致，作者在UperHead [54]框架上使用了一个在ImageNet上预训练的MSVMamba主干网络。训练过程在160K次迭代中进行，批处理大小为16。作者采用AdamW优化器，学习率设置为。作者的实验主要使用默认的输入分辨率进行。此外，作者还融入了多尺度（MS）测试来评估性能变化。
结果。 作者在表3中展示了作者模型以及其他竞争者的详细结果，包括单尺度和多尺度测试。作者的MSVMamba在微小版本上始终优于Swin、ConNeXt和VMamba模型，分别提高了+2.2、+1.6和+0.3 mIoU。
Ablation Study
为了验证所提出模块的有效性，作者进行了一项全面的消融研究。具体来说，作者将VMamba-Tiny模型进行了缩放，将其嵌入维度设置为48，状态空间维度设置为8，以及四个不同阶段的块数设置为。这种缩放模型，作者称之为VMamba-Nano，其参数量和计算成本分别为4.4M和0.87GFLOPs。这个模型作为作者消融实验的 Baseline 。消融研究中的模型在100个周期的训练计划下进行，除非有另外说明，以减少训练时间。
关于多尺度2D扫描。 在保持大致相等FLOP计数的情况下，将VMamba中的SS2D替换为作者的MS2D，准确度从69.6提升到71.9，如表4所示。此外，作者对多尺度扫描中的扫描数量进行了消融研究，考虑了全分辨率和半分辨率分支。结果展示在表5中。
将所有扫描放在半分辨率分支导致显著损失了细粒度特征，从而大幅降低了模型准确度。与只有一个全分辨率分支相比，将两个或三个扫描放在全分辨率分支，准确度分别提高了0.1和0.6，但引入了大约12和25的额外计算成本。将四个扫描分配到全分辨率分支，实际上回到了SS2D方法，计算成本增加了34，但只提高了0.4的准确度。为了在计算成本和准确度之间取得最佳平衡，作者选择全分辨率分支中的一个扫描作为默认设置。在MS2D的基础上，作者引入了EfficientVManba[38]后的SE块，进一步在几乎没有额外计算成本的情况下提高了准确度0.5。
关于ConvFFN。 在用MS2D替换SS2D并引入SE块之后，作者构建了一个使用ConvFFN作为通道混合器的模型。当只使用SSM时，模型在通道间的信息交换不足。将ConvFFN作为通道混合器显著增强了模型进行通道间信息交互的能力。
如表4所示，加入ConvFFN使得准确度又提高了2.0。为了保持大致相等的计算成本，作者调整了模型内的块数。此外，作者将状态空间维度，并堆叠一个额外的块以进一步增强捕捉长距离信息的能力，同时保持大致恒定的计算成本。这一操作使得准确度又提高了0.7，如表4所示。
5 Limitations
多尺度VMamba的设计旨在解决参数有限的Mamba模型在视觉任务中的长距离遗忘问题。尽管提出的模型已被证明是有效的，但其可扩展性仍有待探索，因为通过增加模型大小也可以缓解这一问题。在这种情况下，多尺度设计可能只会带来边际改进。
6 Conclusion
在本文中，作者引入了多尺度VMamba（MSVMamba），这是一种基于SSM的视觉基础模型，它利用了线性复杂性和全局感受野的优势。作者开发了多尺度2D（MS2D）扫描技术，以最小化计算冗余并缓解参数受限视觉模型中的长距离遗忘问题。此外，作者还融入了卷积前馈网络（ConvFFN）以增强通道间信息的交换，从而显著提高了作者模型的性能。作者的实验表明，MSVMamba在图像分类和下游任务中一致性地优于包括ConvNeXt、Swin Transformer和VMamba在内的各种架构的流行模型。
Throughput Comparison
作者在表6中展示了作者MSVMamba变体以及 Baseline VMamba-Tiny在ImageNet-1K数据集上的吞吐量比较。吞吐量是在NVIDIA GeForce RTX 2080 Ti GPU上处理的每秒图像数。所有测试的图像大小为，批量大小设置为128。在相似的GFLOPs下，作者的MSVMamba-Tiny模型几乎将 Baseline VMamba-Tiny的吞吐量翻了一番。
Qualitative Analysis
作者提供了图6中的可视化，比较了VMamba中提出的MS2D和SS2D配置。这些可视化是通过将S6层转换为VMamba [31]展示的关注格式生成的。结果显示，MS2D扫描的全分辨率捕捉到了更多细节特征，而与SS2D相比，半分辨率扫描主要关注更广泛的结构细节。所提出的分层扫描模式增强了当前层在不同抽象层次间辨别和融合特征的能力。
参考
[1].Multi-Scale VMamba: Hierarchy in Hierarchy Visual State Space Model.
点击上方卡片，关注 「AI视界引擎」 公众号