题目：LocalMamba: Visual State Space Model（视觉状态空间模型） with Windowed Selective Scan（窗口化的选择扫描）

论文：[2403.09338] LocalMamba: Visual State Space Model with Windowed Selective Scan (arxiv.org)

源码：hunto/LocalMamba: Code for paper LocalMamba: Visual State Space Model with Windowed Selective Scan (github.com)

目录

一、摘要 

二、引言 

三、方法

3.1 Local Scan for Visual Representations（局部扫描 for 视觉表征）

3.2 搜索自适应扫描

3.3 架构变体

---

一、摘要 

研究背景：状态空间模型的最新进展，特别是Mamba，已经证明了在语言理解等任务的长序列建模方面取得了重大进展。然而，它们在视觉任务中的应用并没有明显超过传统的卷积神经网络（CNNs）和视觉变换器（ViTs）的性能。

研究问题：传统的 Vim 方法展平空间 token，忽略了局部 2D 依赖关系的保持，从而拉长了相邻  token 之间的距离。

主要工作：

• 本文提出了增强 Vision Mamba（Vim）的关键在于优化序列建模的扫描方向。
• 本文提出了一种新颖的 局部扫描策略 ，将图像分割成不同的窗口，在保持全局视角的同时有效地捕捉局部依赖性。
• 此外，考虑到不同网络层对扫描模式的偏好不同，提出了一种动态方法来独立搜索每层的最佳扫描选择，从而显著提高了性能。

实验效果：在普通和分层模型上的广泛实验，强调了所提出方法在有效捕获图像表示方面的优越性。例如，在ImageNet上，我们的模型在相同的1.5G FLOPs下，性能显著优于 Vim-Ti 3.1%。

  

二、引言 

研究现状：SSM 以其计算效率为特征，表现出线性或近线性的序列缩放复杂性长度，使它们特别适合处理长序列。随着 Mamba (一种包含选择性扫描(S6)的新变体) 的成功，SSM 在广泛的视觉任务中的应用激增。

研究问题：图像中 2D 空间模式的非因果性本质上与 SSM 的因果处理框架不一致。如下图所示，将空间数据扁平化为 1D 标记的传统方法破坏了自然的局部 2D 依赖关系，削弱了模型准确解释空间关系的能力。

主要工作：

• 在这项工作中，引入了一种新的方法，通过 将图像分割到多个不同的局部窗口 来改善 Vision Mamba (ViM) 中的局部表示。在进行遍历之前，每个窗口都被单独扫描跨窗口，确保同一2D 语义区域内的标记被紧密地一起处理。
• 通过 整合传统的全局扫描方向和新颖的局部扫描技术 来设计本文的基础块，使模型能够吸收全面的全局和细致入微的局部信息。
• 为了更好地从这些不同的扫描过程中聚合功能，提出了一个 空间和通道注意力模块SCAttn ，旨在识别和强调有价值的信息，同时过滤掉冗余。
• 考虑到扫描方向对特征表示的不同影响(例如，窗口大小为3的局部扫描擅 长捕获较小的物体或细节，而窗口大小为7的更适合于捕获较大的物体)，引入了一种方向搜索方法来选择最佳扫描方向 。这种可变性在不同的层数和网络深度上尤为明显。

主要贡献：

• 为SSM引入了一种新的扫描方法，其中包括在不同窗口内进行 local 扫描，从而显著增强了模型结合全局上下文捕获详细本地信息的能力。
• 引入了一种跨不同网络层搜索扫描方向的方法，能够识别和应用最有效的扫描组合，从而提高网络性能。
• 提出了两个模型的变体，设计与平原和层次结构。通过对图像分类，对象检测和语义分割任务的广泛实验，证明了本文的模型比以前的作品取得了显着的改进。例如，在语义分割任务中，使用类似数量的参数，LocalVimS 在 mIoU (SS) 上的性能比 Vim-S 高出 1.5。

三、方法

LocalMamba整体模型结构图，如下所示： 

3.1 Local Scan for Visual Representations（局部扫描 for 视觉表征）

研究问题：Vim 中使用的扁平化方法破坏了这些局部依赖性，显著增加了垂直相邻令牌之间的距离，并阻碍了模型捕捉局部细微差别的能力。虽然 Vmamba 试图通过在水平和垂直方向上扫描图像来解决这个问题，但它仍然不能在单次扫描中全面处理空间区域。

解决方法：引入了一种局部扫描图像的新方法。通过将图像划分为多个不同的局部窗口，本文方法确保了相关局部 token 的更紧密排列，增强了对局部依赖关系的捕获。

过程：

1）四个方向的选择性扫描机制：通过集成四个方向的选择性扫描机制来构建基础模块：原始(a)和(c)方向，以及它们的翻转（backward 方向），以便从尾部到头部进行扫描。这种多方面的方法确保了每个选择性扫描模块内的全面分析，在局部细节和全局视角之间取得了平衡。

2）SCAttn模块：为了增强不同特征的整合并消除无关信息，在合并前引入了空间和通道关注模块（SCAttn）。该模块自适应地对每个分支的特征内的通道和令牌进行加权。通道注意力分支通过对空间维度上的输入特征 求平均值来聚合全局表示，随后应用线性变换来确定通道权重 。相反，空间注意力机制通过 用全局表示来增强每个标记的特征来评估标记意义，从而实现细致入微的重要性加权特征提取。（通道注意力 + 全局投票门）

3.2 搜索自适应扫描

问题：四个分支局部选择性扫描模块，这种方法大大增加了计算需求。

解决方法：引入了一种策略，可以有效地为每层选择最合适的扫描方向，从而在不产生过多计算成本的情况下优化性能。该方法涉及为每层搜索最优扫描配置，确保定制和高效的表示建模。

过程：为了定制每一层的扫描过程，引入了 8 个候选扫描方向的多样化集合 S。这些包括水平和垂直扫描 (标准和翻转)，以及窗口大小为 2 和 7 的局部扫描。为了与以前的模型保持一致的计算预算， 从这 8 个方向中为每层选择 4 个方向。这种方法的结果是一个可观的搜索空间 ，其中 K 表示模块的总数。

原理：将离散选择过程转换为连续域，允许使用softmax概率来表示扫描方向的选择：

其中， 表示每个层  的一组可学习参数，反映所有潜在扫描方向上的 softmax 概率。（通过连续域中的可学习参数进行 Softmax 函数激活并来连续求和，找出最高softmax概率的方向组合）

（将整个搜索空间构建为一个过度参数化的网络，能够同时优化网络参数和架构变量 α，遵循标准的训练协议。在训练完成后，通过选择具有最高softmax概率的四个方向来得出最优方向选项。）

模型的搜索方向可视化结果，如下图所示：

3.3 架构变体

主干：作者尝试了两种不同的主干 ViM 和 VMamba。

实施细节：原始 Vim 模块包括两个扫描方向（水平和翻转水平），本文为LocalVim引入了四个扫描方向，从而增加了计算开销。为了保持类似的计算预算，我们将 Vim 模块的数量从 24 调整为20。对于LocalVMamba，它本质上有四个扫描方向，类似于本文提出的模块，直接替换模块，而不改变结构配置。详细配置可见下表：

计算成本分析：LocalMamba 模块是高效和有效的，只增加了很小的计算成本。扫描机制只涉及重新定位标记，不会产生额外的计算开销。此外，SCAttn模块，旨在跨扫描有效聚合各种信息，是异常精简。它利用线性层将token维度减少 1/r 的因子，随后在空间和通道维度上生成注意力权重，所有模型的 r 设置为8。例如，LocalVMamba-T模型，用LocalMamba模块替换了 VMamba模块，只会将VMamva-T的FLOPs从5.6G增加到5.7G。