商汤&悉尼大学提出LocalMamba | 局部扫描机制让Mamba保留全局视角同时捕获局部信息

近期在状态空间模型方面，尤其是Mamba的研究，已经在诸如语言理解等任务中对长序列建模方面取得了显著进展。然而，在视觉任务中的应用并没有明显超越传统的卷积神经网络（CNNs）和视觉 Transformer （ViTs）。

本文提出，提高视觉Mamba（ViM）性能的关键在于优化序列建模的扫描方向。传统的ViM方法通过展平空间标记，忽视了局部2D依赖关系的保持，从而增加了相邻标记之间的距离。作者引入了一种新颖的局部扫描策略，将图像划分为不同的窗口，有效捕获局部依赖关系的同时保持了全局视角。

此外，考虑到不同网络层对扫描模式的不同偏好，作者提出了一种动态方法，独立地为每一层搜索最优的扫描选择，极大地提升了性能。

在普通和分层模型上的大量实验强调了LocalVim在有效捕捉图像表示方面的优越性。例如，在相同的1.5G FLOPs下，LocalVim在ImageNet上的性能比Vim-Ti显著高出3.1%。

代码：https://github.com/hunto/LocalVim

1 Introduction

结构化状态空间模型（SSMs）最近在序列建模领域作为一种多功能的架构受到重视，预示着在计算效率和模型多样性之间达到平衡的新时代的到来。这些模型综合了循环神经网络（RNNs）和卷积神经网络（CNNs）的最佳属性，并从经典状态空间模型的基本原理中汲取灵感。以其计算效率为特点，SSMs展现出与序列长度成线性或近线性复杂度的扩展性，使它们特别适合处理长序列。

在Mamba成功的基础上，一种融合了选择性扫描（S6）的新变体出现，此后将SSMs应用于广泛的视觉任务中便有了激增。这些应用从开发通用基础模型到推动图像分割和合成等领域的发展，展示了模型在视觉领域的适应性和潜力。

通常这些视觉研究需要将2D图像转换成1D序列以进行基于SSM的处理，然后将Mamba的原始SSM结构整合到他们针对特定任务的基础模型中。然而，它们仅比传统的卷积神经网络（CNNs）和视觉 Transformer （ViTs）显示出适度的改进。这种适度的进步凸显了一个重大挑战：图像中2D空间模式的非因果关系与SSMs的因果处理框架本质上是冲突的。

如图1所示，传统方法将空间数据扁平化为1D标记，破坏了自然的局部2D依赖关系，削弱了模型准确解释空间关系的能力。尽管V Mamba引入了一种2D扫描技术，通过在水平和垂直方向上扫描图像来解决这个问题，但它仍然难以在扫描序列中保持原本相邻标记的邻近性，这对于有效的局部表示建模至关重要。

在这项工作中，作者提出了一种新颖的方法来提高视觉Mamba（ViM）中的局部表示，通过将图像分割成多个独特的局部窗口。在跨窗口进行遍历之前，每个窗口都单独进行扫描，确保同一2D语义区域内的标记被紧密处理。这种方法显著提高了模型捕捉局部区域之间细节的能力，实验结果在图2中得到验证。

作者设计的基础块集成了传统的全局扫描方向和作者的新型局部扫描技术，使模型能够融合全面的全局信息和细致的局部信息。此外，为了更好地聚合来自这些不同扫描过程的特点，作者提出了一个空间和通道注意力模块，SCAttn，旨在辨别并强调有价值的信息同时过滤掉冗余。

承认扫描方向对特征表征的不同影响（例如，窗口大小为3的局部扫描擅长捕捉较小目标或细节，而窗口大小为7则更适合较大目标），作者引入了一种方向搜索方法以选择最优的扫描方向。这种变异性在不同层和网络深度中尤为明显。受到DARTS的启发，作者从离散选择推进到连续域，用一个可学习的因子来表示，以在单一网络内融合多个扫描方向。在这个网络训练完成后，通过识别具有最高分配概率的方向来确定最有效的扫描方向。

作者开发的模型，LocalVim和LocalVAMaba，融合了平面和层次化的结构，这使得相较于之前的方法，作者取得了显著的提升。

本研究的主要贡献包括：

1. 作者引入了一种新颖的SSM扫描方法，包括在各自窗口内进行局部扫描，这显著提高了LocalVim在结合全局上下文的同时捕捉详细局部信息的能力。
2. 作者开发了一种方法，用于搜索不同网络层中的扫描方向，这使得作者能够识别并应用最有效的扫描组合，从而提高网络性能。
3. 作者提出了两种模型变体，分别采用普通和分层结构设计。通过对图像分类、目标检测和语义分割任务的广泛实验，作者展示了LocalVim相较于之前的工作取得了显著改进。例如，在语义分割任务上，在参数量相似的情况下，LocalVim-S模型在mIoU (SS)上以1.5的大幅度优于Vim-S模型。

2 Related Work

Generic Vision Backbone Design

过去十年见证了计算机视觉领域的重大变革性进步，这主要是由深度神经网络的演变和基础通用模型的涌现所推动的。最初，卷积神经网络（CNNs）在视觉模型架构上树立了一个重要的里程碑，为复杂的图像识别和分析任务奠定了基础。

在这些工作中，ResNet凭借其基石残差连接技术，是最受欢迎的模型之一，被广泛应用于广泛的视觉任务领域；MobileNet系列通过使用深度可分离卷积，引领了轻量级模型的设计。然而，Vision Transformer (ViT) 的引入标志着一个范式的转变，挑战了CNN在该领域的霸主地位。

ViT通过将图像分割成一系列顺序的块，并利用Transformer架构的核心组件——自注意力机制来提取特征，从而革新了图像处理的方法。这种新颖的方法突显了Transformer在视觉任务中尚未发掘的潜力，激发了旨在改进其架构设计和训练方法的研究热潮，提高了计算效率，并扩展了其应用范围。

在Mamba这一状态空间模型（SSMs）变体的长序列建模成功基础上，一些创新模型如Vim和V Mamba被引入到视觉任务中，即Vision Mamba。这些模型适应Mamba框架，作为一个通用的视觉应用 Backbone 网络，展示了在高分辨率图像上相较于传统CNN和ViT的卓越效率和准确度。

State Space Models

状态空间模型（SSMs）代表了一种为序列到序列转换设计的架构范式，擅长处理长依赖标记。尽管在训练初期由于计算和内存的强度问题面临挑战，但近期进展已显著改善这些问题，使得深度SSMs成为了与CNNs和Transformers相抗衡的强大竞争者。

特别是，S4引入了一种高效的正常加低秩（NPLR）表示，利用Woodbury恒等式进行矩阵求逆加速，从而简化了卷积核计算。在此基础上，Mamba通过结合特定输入参数化及可扩展、硬件优化计算方法进一步精化了SSMs，在处理跨语言和基因组学的广泛序列时，实现了前所未有的效率和简洁性。

S4ND的出现标志着SSM块在视觉任务中的首次尝试，它能熟练地在1D、2D和3D领域将视觉数据作为连续信号进行处理。随后，受到Mamba模型的成功的启发，Vmamba和 Vim拓展到通用视觉任务中，通过提出双向扫描和交叉扫描机制来解决SSMs中的方向敏感性挑战。

在Mamba在通用模型基础之上，为视觉任务开发了新的方法，如图像分割和图像合成，展示了视觉Mamba模型在应对复杂的视觉挑战时的适应性和有效性。

Preliminaries

State Space Models

结构化状态空间模型（SSMs）是深度学习中一类序列模型的代表，其特点在于能够将一维序列 通过中间潜在状态 映射到 。

在该系统中，系统矩阵 , 和 分别控制动态和输出映射。

离散化。 在实际实施中，通过使用零阶保持假设4，将方程1描述的连续系统离散化，有效地将连续时间参数（ ， ）转换为在特定采样时间尺度 上的离散对应参数（ ， ）：

这导致了一个离散化模型的公式如下：

为了计算效率，可以通过并行计算加快方程3中描述的迭代过程，采用全局卷积操作：

其中 表示卷积操作，而 作为SSM的核。这种方法利用卷积同时合成序列中的输出，提高了计算效率和可扩展性。

Selective State Space Models

传统的状态空间模型（SSMs），通常指的是S4，已实现了线性时间复杂度。然而，它们捕捉序列上下文的能力本质上受到静态参数化的限制。

为了解决这一限制，选择性的状态空间模型（称为Mamba）引入了一种动态和选择性的机制来处理序列状态之间的交互。与使用恒定转换参数 的传统SSMs不同，Mamba模型采用输入依赖的参数，使得模型能够拥有更丰富、对序列敏感的参数化。

具体来说，Mamba模型直接从输入序列 计算参数 ， ，以及 。

Mamba模型利用选择性的SSMs，不仅实现了序列长度的线性可扩展性，而且在语言建模任务中也提供了有竞争力的性能。这一成功激发了后续在视觉任务中的应用，研究提出了将Mamba整合到基础视觉模型中。Vim采用了类似ViT的架构，用双向Mamba块取代了传统的 Transformer 块。

V Mamba引入了一种新颖的2D选择性扫描技术，以水平和垂直方向扫描图像，并构建了一个类似于Swin Transformer的分层模型。作者的研究扩展了这些初步探索，专注于为视觉任务优化S6适配，在作者这里取得了性能提升的结果。

4 Methodology

本部分详细阐述了LocalVim的核心组件，首先是从图像中增强模型挖掘细粒度细节能力的本地扫描机制。随后，作者介绍了扫描方向搜索算法，这是一种创新方法，它能识别不同层间的最优扫描序列，从而确保全局和局部视觉线索的和谐整合。本节的最后一部分展示了LocalVim框架在简单平面架构和复杂分层架构中的部署，展示了其在多样化设置中的灵活性和有效性。

Local Scan for Visual Representations

作者的方法采用了选择扫描机制S6，该机制在处理一维因果序列数据时表现出色。这种机制因果地处理输入，有效地捕捉到扫描段内的关键信息，类似于语言建模中理解序列词之间的依赖关系至关重要。

然而，图像中二维空间数据的固有非因果性质，对这种因果处理方法提出了重大挑战。传统的将空间标记扁平化的策略破坏了局部二维依赖的完整性，从而降低了模型有效辨别空间关系的能力。

例如，如图1（a）和（b）所示，Vim中使用的扁平化方法破坏了这些局部依赖，显著增加了垂直相邻标记之间的距离，阻碍了模型捕捉局部细微差别的能力。尽管VMDA尝试通过水平和垂直两个方向扫描图像来解决这一问题，但它仍然未能在一个扫描过程中全面处理空间区域。

为了解决这一限制，作者引入了一种新颖的方法用于局部扫描图像。通过将图像划分为多个独特的局部窗口，作者的方法确保了相关局部标记的更紧密排列，从而增强了对局部依赖性的捕捉。这项技术在图1（c）中有所展示，与那些无法保持空间连贯性的先前方法进行了对比。

尽管作者的方法在有效捕捉每个区域内的局部依赖性方面表现出色，但它也认识到全局上下文的重要性。为此，作者通过在四个方向上整合选择扫描机制来构建作者的基础模块：原始的（a）和（c）方向，以及它们的翻转方向，后者支持从头到尾的扫描（在Vim和VMDA中采用翻转方向，以便更好地建模非因果图像标记）。这种多方面的方法确保了在每个选择扫描模块中进行全面分析，同时在局部细节和全局视角之间保持平衡。

如图3所示，作者的模块通过四个独特的选择性扫描分支处理每个输入图像特征。这些分支独立捕捉相关信息，随后将这些信息合并成一个统一特征输出。为了增强不同特征的综合集成并消除无关信息，作者在合并之前引入了一个空间和通道注意力模块。

如图2(b)所示，该模块自适应地对每个分支特征中的通道和标记进行加权，包括两个关键组成部分：一个通道注意力分支和一个空间注意力分支。通道注意力分支通过在空间维度上平均输入特征来聚集全局表示，然后应用线性变换来确定通道权重。相反，空间注意力机制通过将每个标记的特征与全局表示相结合，评估标记的显著性，从而实现细致的重要性加权特征提取。

备注。 尽管一些ViT变体，例如Swin Transformer，提出了将图像划分为更小窗口的方法，但作者的LocalVim中的局部扫描在目的和效果上都是独特的。ViTs中的窗口自注意力主要是为了解决全局自注意力的计算效率问题，尽管这牺牲了一些全局注意力能力。

相反，局部扫描机制旨在重新排列标记位置，以增强视觉Mamba中对局部区域依赖关系的建模，同时仍保留全局理解能力，因为整个图像仍然通过SSM进行聚合和处理。

Searching for Adaptive Scan

结构化状态空间模型（SSM）在捕捉图像表示方面的有效性因扫描方向的不同而有所差异。直观上，要达到最优性能，需要在不同方向上进行多次扫描，类似于作者之前讨论的4分支局部选择扫描块。然而，这种方法大大增加了计算需求。

为了解决这个问题，作者提出了一种策略，可以有效地为每一层选择最合适的扫描方向，从而在不过度增加计算成本的情况下优化性能。这种方法涉及为每一层搜索最优的扫描配置，确保进行定制且高效的表示建模。

搜索空间。 为了针对每一层定制扫描过程，作者引入了一个多样化的候选扫描方向集合 ，包含8种方向。这些方向包括水平扫描和垂直扫描（包括标准和翻转两种），以及窗口大小为2和7的局部扫描（同样也包括标准和翻转两种）。

为了保持与之前模型一致的计算预算，作者为每一层从这8个方向中选择4个。这种方法导致了一个相当大的搜索空间，即 ，其中 表示总块数。

基于DARTS的原则，作者的方法应用了一种可微分的搜索机制来寻找扫描方向，并采用连续松弛法来导航分类选择。这种方法将离散的选择过程转换为连续的领域，允许使用softmax概率来表示扫描方向的选择：

其中 表示每一层 的一组可学习参数，反映了所有潜在扫描方向上的softmax概率。

作者构建了整个搜索空间作为一个过度参数化的网络，这使作者能够按照标准的训练协议同时优化网络参数和架构变量 。在训练完成后，作者通过选择具有最高softmax概率的四个方向来推导出最佳的方向选项。

作者在图4中可视化了作者模型的搜索方向。有关搜索结果的详细分析，请参见第5.5节。

方向搜索的可扩展性。 作者当前的方法在训练中汇总了所有扫描方向以供选择，恰好适用于具有适度选项范围的模式。例如，一个具有20个模块和每个模块128个方向的模型需要28 GB的GPU内存，这表明在广泛选择时存在可扩展性的限制。

为了在具有大量选择阵列的场景中减轻内存消耗，诸如单路径采样，二进制近似，以及部分通道使用 [55] 等技术提供可行的解决方案。作者将对更自适应的方向策略和先进搜索技术的研究留待未来的努力。

Architecture Variants

为了彻底评估作者的方法的有效性，作者引入了基于普通和分层结构的架构变体，分别命名为LocalVim和LocalVamba。这些架构的配置在表1中详细说明。特别是，在LocalVim中，标准的SSM块被LocalVim块所替代，如图3所示。

考虑到原始Vim块包括两种扫描方向（水平和翻转水平），而LocalVim引入了四种扫描方向，从而增加了计算开销。为了保持类似的计算预算，作者将Vim块的数量从24调整到20。对于固有就有四种扫描方向的LocalVMDA，作者直接替换块而不改变结构配置。

计算成本分析。 作者的LocalVim模块既高效又有效，仅在计算成本上有很小的增加。扫描机制仅涉及重新定位token，因此在FLOPs上不产生额外的计算成本。

此外，SCAttn模块旨在高效地跨扫描聚合不同的信息，其设计特别简化。它利用线性层将token维度减少到原来的 ，然后在空间和通道维度上生成注意力权重，其中 对所有模型设置为8。例如，LocalVMDA-T模型，用LocalVim块替换了VMDA块，仅将Vdamva-T的FLOPs从5.6G增加到5.7G。

5 Experiments

本节概述了作者的实验评估，首先从ImageNet分类任务开始，随后将训练好的模型迁移到各种下游任务中，包括目标检测和语义分割。

ImageNet Classification

训练策略。 作者在ImageNet-1K数据集上训练模型，并在ImageNet-1K验证集上评估性能。与之前的工作一致，作者使用基础批处理大小为1024和AdamW优化器训练模型300个周期，采用余弦退火学习率计划，初始值为 ，并设有20个周期的 Warm up 期。

对于训练数据增强，作者使用随机裁剪，AutoAugment策略为rand-m9-mstd0.5，以及在每个图像上以0.25的概率随机擦除像素，然后在每个批次中采用MixUp策略，其比率为0.2。模型上采用指数移动平均，衰减率为0.9999。

扫描方向搜索。 对于超网络训练，作者将训练周期减少到100次，同时保持其他超参数与标准的ImageNet训练一致。在LocalVim变体中，超网络的嵌入维度设置为128，由于LocalVim-T和LocalVim-S具有统一的层结构，因此对它们的搜索操作是相同的。对于LocalV-Mamba变体，包括LocalVMDA-T和LocalVMDA-S，初始嵌入维度被最小化到32，以促进搜索过程。

结果。 如表2所示，作者的结果显示，相较于传统的CNN和ViT方法，准确度有了显著的提升。特别值得一提的是，LocalVim-T在1.5G FLOPs的情况下达到了76.2%的准确率，超过了准确率为72.2%的DeiT-Ti。

在分层结构中，LocalVMDA-T的82.7%准确率比Swin-T高出1.4%。此外，与作者的seminal贡献，Vim和V Mamba相比，作者的方法取得了显著的增益；例如，LocalVim-T和LocalV Mamba-T在准确度上分别超过Vim-Ti和V Mamba-T 2.7%和0.5%。

另外，为了验证局部扫描的有效性，作者在没有扫描方向搜索的模型上进行了额外的实验，这些内容在4.1节中有详细描述，并在表中以记。仅将局部扫描融入到原始Vim框架中，LocalVim-T 就比Vim-Ti的准确率高出了2.7%，而完整的方法进一步将准确度提高了0.4%。这些发现确认了扫描方向在视觉SSMs中的关键作用，证明了作者局部扫描方法能有效增强局部依赖的捕获能力。

Object Detection

训练策略。 作者使用MSCOCO 2017数据集和MMDetection库来验证目标检测的性能。对于LocalV Mamba系列遵循以前的工作，使用Mask-RCNN检测器来训练目标检测和实例分割任务。训练策略包括1 设置，即12个训练周期，以及3 设置，即36个训练周期，并采用多尺度数据增强。而对于LocalVim，作者遵循Vim的方法，使用级联Mask R-CNN结合ViTDet作为检测器。

结果。作者在表3中总结了LocalVAMaba与其他基础网络的对比结果。可以看出，作者的LocalVAMaba在所有模型变体上均一致优于VAMaba。

与其他架构，即卷积神经网络（CNNs）和视觉 Transformer （ViTs）相比，作者获得了显著的优势。例如，LocalVAMaba-T获得了46.7的边界框AP和42.2的 Mask AP，分别比Swin-T提高了4.0和2.9的大幅差距。关于与Vim的定量比较，请参考补充材料。

Semantic Segmentation

训练策略。在ADE20K数据集上用作者的 Backbone 网络训练UperNet。模型以总批次大小16和 的输入进行训练，采用了带有权重衰减0.01的AdamW优化器。作者使用多项式学习率调度，经过160K次迭代衰减，初始学习率为 。注意，Vim没有报告FLOPs和mIoU（MS），也没有发布分割的代码，因此作者按照MMSegmentation中的ViT示例配置实现了LocalVim。

结果。作者在表4中报告了LocalVim和LocalVMamba的结果。在LocalVim上，相比于 Baseline 模型Vim-Ti，作者取得了显著的改进。例如，在相似的参数量下，LocalVim-S在mIoU（SS）上比Vim-S高出1.5个点。

而在LocalVMamba上，相比于VMamba Baseline ，作者也取得了显著的改进；例如，作者的LocalVMamba-T达到了49.1的出色mIoU（MS）分数，比VMamba-T高出0.8。与CNN和ViT相比，作者的改进更为明显。这些结果证明了SSMs全局表示在密集预测任务中的有效性。

Ablation Study

局部扫描的影响。本文评估了作者的局部扫描技术的影响，具体实验细节见表5。将Vim-T的传统水平扫描替换为作者的局部扫描，相较于 Baseline 方法，性能提升了1%。

在受限的FLOP预算下，LocalVim-T 中扫描方向的组合又带来了额外的1.1%准确度提升。这些结果强调了不同窗口大小（将水平扫描视为窗口大小为 的局部扫描）在图像识别中的不同影响，而这些扫描的综合运用进一步提升了性能。

SCAttn的效果。在表5中，将SCAttn整合到最后的LocalVim块中，实现了额外的0.6%的提升，验证了策略性地结合不同扫描方向的有效性。这强调了SCAttn通过自适应地合并扫描方向来提高性能的作用。

扫描方向搜索的影响。 如表2所示，作者的实证评估证实了扫描方向搜索策略在最终的LocalVim模型中带来的显著好处。这些模型相比于仅合并水平扫描、窗口大小为 的局部扫描及其镜像的版本，显示出明显的改进。

例如，LocalVim-T相较于LocalVim-T 提高了 。这种性能提升可以归因于每一层扫描组合的方法论选择，这为优化模型有效性提供了一套多样化的选项。

Visualization of Searched Scan Directions

图4展示了作者模型中获得的扫描方向的可视化。观察结果表明，在LocalVim的简单架构内部，倾向于在初始和最终段使用局部扫描，而中间层则偏好全局水平和垂直扫描。

值得注意的是， 局部扫描倾向于集中在网络的尾部，而更大的 扫描在网络初始部分更为突出。相反，LocalVAMaba的分层结构相比于LocalVim更倾向于使用局部扫描，且更偏好 扫描而不是 扫描。

6 Conclusion

在本文中，作者介绍了LocalVim，这是一种视觉状态空间模型的新方法，它显著增强了图像内部局部依赖关系的捕捉能力，同时保持了全局上下文理解。作者的方法利用了窗口化选择性扫描和扫描方向搜索，从而显著优于现有模型。

在多个数据集和任务上的广泛实验已经证明了LocalVim相较于传统卷积神经网络（CNN）和视觉 Transformer （ViTs）的优越性，为图像分类、目标检测和语义分割建立了新的基准。作者的研究强调了扫描机制在视觉状态空间模型中的重要性，并为高效和有效的状态空间建模研究开辟了新的途径。未来的工作将探索作者的方法扩展到更复杂和多样化视觉任务的可行性，以及整合更先进的扫描策略的潜力。

Comparison to Vim on Object Detection

与使用Mask R-CNN框架进行基准测试的V Mamba不同，Vim采用了ViTDet中的neck架构，并以Cascade Mask R-CNN作为检测器进行训练。作者也与Vim的设置保持一致，以便对LocalVim模型进行公平的比较和评估。

作者在表6中总结了结果。LocalVim-T与Vim-Ti表现相当，但在AP 和mask AP方面具有显著优势。例如，LocalVim-T在AP 和AP 上分别比Vim-Ti提高了0.7和2.1。

Visualization of Searched Directions on LocalVMamba-S

作者在图5中可视化了LocalVMamba-S搜索的方向。在这个模型中，在第3阶段有27层的情况下，相比于LocalVMamba-T，更偏好使用 的局部扫描。

Discussions

潜在的负面影响。 研究所提出模型的效应需要大量的计算资源消耗，这可能会潜在地引发对环境问题的关注。

局限性。 视觉状态空间模型具有与序列长度成线性复杂度，与先前的CNN和ViTs架构相比，在大型分辨率下游任务上显示出显著的改进。然而，SSM的计算框架本质上比卷积和自注意力机制更为复杂，这增加了并行计算的效率执行难度。

当前的深度学习框架在加速SSM计算方面也表现出比支持更成熟架构时要有限的能力。积极的一面是，诸如VMamba5等项目中正在进行的工作旨在提高选择性SSM操作的计算效率。这些努力已经在速度上实现了显著的进步，正如在Mamba的原始实现中所记录的改进所证明的那样。

参考

[1].LocalVim: Visual State Space Model with Windowed Selective Scan.