LocalMamba

标题：《LocalMamba: Visual State Space Model with Windowed Selective Scan》
论文：https://arxiv.org/pdf/2403.09338.pdf
源码：https://github.com/hunto/LocalMamba

导读

近期，以 Mamba 为代表的状态空间模型（State Space Models）在诸如自然语言理解等序列建模任务中取得了显著进展。然而，其在视觉任务中的应用并没有表现出相较卷积神经网络（CNNs）和视觉 Transformer（ViTs）的明显优势。

来自悉尼大学、商汤等机构的研究者指出，不同于视觉 Transformer，Mamba中序列的扫描方向对其精度提升至关重要。传统的视觉Mamba（ViMs）方法只是简单地将图像token通过逐行或逐列扫描铺平成一维序列，从而增加了图像中相邻token的距离，弱化了局部2D依赖关系的建模。

因此，作者引入了一种新颖的局部扫描策略，通过将图像划分为不同的窗口，有效地捕捉了局部依赖关系，同时保持Mamba原有的全局依赖能力。此外，作者还提出了一种动态搜索方法，用于为不同层搜索其最优的扫描方向组合，显著提升了性能。此方法得到的LocalMamba模型，刷新了Mamba模型在通用视觉任务上新的SOTA。代码已开源，欢迎 star！

动机

状态空间模型（SSMs）凭借其对序列长度具有线性时间复杂度的优势，在自然语言处理等序列建模任务中表现出了更好的计算效率与精度平衡。在Mamba成功的基础上，诸如Vim [2]、VMamba [3]等方法将SSM应用到了通用视觉任务上，并在图像分辨率较高的场景下取得了显著的计算量及内存占用优势。

然而，如图1所示，先前的方法逐行或逐列扫描图像，造成了原本在2D空间上相邻的token的割裂，增加了局部区域特征提取的难度。为了解决这一问题，本文提出了一种局部扫描机制用于提升视觉Mamba的局部表示能力。该机制首先将图片划分成多个独立窗口，再使用先窗口内扫描、后窗口间扫描的顺序，保证了属于同一个2D语义空间的token能够彼此靠近地输入SSM进行特征建模。这一机制能够显著地提升模型的局部细节捕捉能力。如图2所示，仅通过给Vim中原有的全局扫描增加局部扫描机制，本方法在相近的计算量下取得了显著的ImageNet分类精度提升。

同时，考虑到不同扫描方向对特征有不同的影响（例如，窗口大小为3的局部扫描擅长捕捉较小的目标和细节，而窗口大小为7则更适合较大目标），本文引入了一种扫描方向搜索方法，以为每层网络选择最优的方向组合。

最终，为了验证方法的有效性，作者在扁平结构的Vim以及层次结构的VMamba上均加入了LocalMamba方法，并在图像分类、目标检测、语义分割任务上进行了实验，相较之前方法取得了明显提升。

方法

局部扫描与全局扫描结合的多分支结构

本文通过实验发现，本文提出的局部扫描机制和原始的全局扫描机制对特征提取有不同的贡献，将二者结合使用能够最大化模型性能。因此，如图2所示，作者使用了一种多分支扫描结构，将图像输入扫描成4个方向（分别为逐行扫描、反向逐行扫描、2x2局部扫描、反向2x2局部扫描），并将SSM处理后的序列融合成一个图像。

为了增强不同扫描特征的融合并消除特征中的无关信息，本文在特征合并前使用了一个注意力模块，用于自适应地根据各个分支特征进行特征加权，最后通过相加的方式将多个分支融合成一个分支。

与视觉Transformer中窗口化注意力的对比。 尽管如Swin Transformer这样的ViT模型也提出了将图像划分为窗口的方法，但是本文得局部扫描机制在目的和效果上都是独特的。ViTs中的窗口自注意力主要是为了解决原始Transformer中全局注意力的计算效率问题，并牺牲了一定的全局注意力能力。而本文得局部扫描机制旨在重新排列图像token位置，以增强视觉Mamba模型对局部依赖的建模能力，同时因为整个图像序列仍然通过SSM进行处理，仍保留了模型的全局理解能力。

自适应的扫描方向搜索

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

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

搜索空间：对于模型的每一层，本文设计了8个候选方向：逐行扫描、反向逐行扫描、逐列扫描、反向逐列扫描、2x2局部扫描、反向2x2局部扫描、7x7局部扫描、反向7x7局部扫描。为了保持与之前模型一致的计算量，每一层从这8个方向中选择4个方向得到最终的扫描模块。

基于可微分模型结构搜索（DARTS）的原则，作者的方法应用了一种可微分的搜索机制来寻找扫描方向，并将离散的选择松弛为连续的softmax概率：

使用这种不同扫描方向的加权求和，本文构建了一个包含所有分支的网络用于搜索方向。通过在原始任务上训练该模型，模型会自适应地学习到对不同方向的偏好，即概率越高代表模型在该层中越倾向于选择这一方向。最终，作者通过选择概率最高的4个方向作为搜索结果。搜索到的方向组合如图4所示。

实验

图像分类

本文使用ImageNet-1K数据集进行图像分类基准测试。如表1所示，作者的结果显示，相较于传统的CNN和ViT方法，LocalMamba准确度有了显著的提升。例如，LocalVim-T在1.5G FLOPs的情况下达到了76.2%的准确率，超过了准确率为72.2%的DeiT-Ti。在分层结构中，LocalVMamba-T的82.7%准确率比Swin-T高出1.4%。此外，与Vim和VMamba相比，作者的方法取得了显著的增益。例如，LocalVim-T和LocalVMamba-T在准确度上分别超过Vim-Ti和VMamba-T 2.7%和0.5%。

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

目标检测与分割

作者在表2中总结了LocalVMamba与其他基础网络的对比结果。可以看出，LocalVMamba在所有模型变体上均一致优于VMamba。与其他架构，即卷积神经网络（CNNs）和视觉 Transformer（ViTs）相比，模型获得了显著的优势。例如，LocalVAMaba-T获得了46.7的检测AP和42.2的分割AP，分别较Swin-T大幅提高了4.0和2.9。

语义分割

作者在表3中报告了LocalVim和LocalVMamba的语义分割结果。在LocalVim上，相比于 基线模型Vim-Ti，作者取得了显著的改进。例如，在相似的参数量下，LocalVim-S在mIoU（SS）上比Vim-S高出1.5个点。而在LocalVMamba上，相比于VMamba，作者也取得了显著的改进。例如，LocalVMamba-T达到了49.1的出色mIoU（MS）分数，比VMamba-T高出0.8。与CNN和ViT相比，作者的改进更为明显。这些结果证明了SSMs全局表示在密集预测任务中的有效性。

总结

在本文中，作者介绍了LocalMamba，一种新的视觉状态空间模型，它显著增强了对图像局部依赖关系的捕捉能力，同时保持了全局上下文理解。作者的方法利用了窗口化选择性扫描和扫描方向搜索，取得了显著优于现有模型的效果。在多个数据集和任务上的广泛实验已经证明了LocalMamba相较于传统卷积神经网络（CNN）和视觉 Transformer（ViTs）的优越性，为图像分类、目标检测和语义分割建立了新的基准。作者的研究强调了扫描机制在视觉状态空间模型中的重要性，并为高效和有效的状态空间建模研究开辟了新的途径。未来的工作将探索作者的方法扩展到更复杂和多样化视觉任务的可行性，以及整合更先进的扫描策略的潜力。

参考文献

[1] Huang, Tao, et al. “LocalMamba: Visual State Space Model with Windowed Selective Scan.” arXiv preprint arXiv:2403.09338 (2024).
[2] Zhu, Lianghui, et al. “Vision mamba: Efficient visual representation learning with bidirectional state space model.” arXiv preprint arXiv:2401.09417 (2024).
[3] Liu, Yue, et al. “Vmamba: Visual state space model.” arXiv preprint arXiv:2401.10166 (2024).