LocalMamba: Visual State Space Model with Windowed Selective Scan 论文总结

向岸看

于 2024-09-15 22:06:03 发布

阅读量886

点赞数 21

文章标签：人工智能深度学习

本文链接：https://blog.csdn.net/qq_45981086/article/details/142214314

版权

题目：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 个方向。这种方法的结果是一个可观的搜索空间 $(C_8^4)^K$ ，其中 K 表示模块的总数。

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

其中， $\alpha^{(l)}$ 表示每个层 $l$ 的一组可学习参数，反映所有潜在扫描方向上的 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。