VMamba: Visual State Space Model

公和众和号：EDPJ（进 Q 交流群：922230617 或加 VX：CV_EDPJ 进 V 交流群）

目录

0. 摘要

3. 方法

3.1 基础概念

3.2 2D 选择性扫描

3.3 VMamba 模型

3.3.1 整体架构

3.3.2 VSS块

4. 实验

---

0. 摘要

卷积神经网络（CNNs）和视觉变压器（Vision Transformers，ViTs）是视觉表示学习的两个最受欢迎的基础模型。虽然 CNNs 在图像分辨率方面具有线性复杂性的显著可扩展性，ViTs 却在拟合能力上超越它们，尽管面临着二次复杂性的挑战。仔细检查发现，ViTs 通过整合全局感受野和动态权重实现了卓越的视觉建模性能。这一观察激发了我们提出一种新颖的架构，该架构继承了这些组件同时增强了计算效率。为此，我们从最近引入的状态空间模型（state space model，SSM）中汲取灵感，提出了视觉状态空间模型（Visual State Space Model，表示为 VMamba），它在不牺牲全局感受野的情况下实现了线性复杂性。为了解决遇到的方向敏感问题，我们引入了交叉扫描模块（Cross-Scan Module，CSM）来遍历空间域并将任何非因果视觉图像转换为有序补丁序列。广泛的实验证明，VMamba 不仅在各种视觉感知任务中表现出有希望的能力，而且在图像分辨率增加时还表现出对已建立基准的更为显著的优势。

源代码：https://github.com/MzeroMiko/VMamba

VMamba 成功有效降低注意力复杂性的关键概念继承自选择性扫描空间状态顺序模型（Selective Scan Space State Sequential Model，S6）[12]，该模型最初设计用于处理自然语言处理（NLP）任务。与传统的注意力计算方法不同，S6 使得 1-D 数组中的每个元素（例如文本序列）能够通过压缩的隐藏状态与先前扫描的任何样本进行交互，有效地将二次复杂性降低到线性。

然而，由于视觉数据的非因果性质，直接将这种策略应用于补丁化和展平的图像将不可避免地导致受限的感受野，因为无法估计相对于未扫描的补丁的关系。我们将这个问题称为 “方向敏感” 问题，并提出通过新引入的交叉扫描模块（Cross-Scan Module，CSM）来解决它。CSM 不是以单向模式（列向或行向）遍历图像特征映射的空间域，而是采用四向扫描策略，即从特征映射的四个角到相对位置（见图 2（b））。这种策略确保特征映射中的每个元素从不同方向的所有其他位置集成信息，从而产生全局感受野，而不增加线性计算复杂性。 

3. 方法

3.1 基础概念

状态空间模型。状态空间模型（SSMs）通常被认为是将刺激 x(t) ∈ R^L 映射到响应 y(t) ∈ R^L 的线性时不变系统。从数学上讲，这些模型通常被构建为线性常微分方程（ODEs）（方程1），其中参数包括 A ∈ C^(N×N)，B、C ∈ C^N，N 为状态大小，以及跳跃连接 D ∈ C^1。

离散化。作为连续时间模型，状态空间模型（SSMs）在集成到深度学习算法中时面临巨大的挑战。为了克服这一障碍，离散化过程变得至关重要。

离散化的主要目标是将常微分方程（ODE）转换为离散函数。这种转换对齐模型与输入数据中所体现的底层信号的采样率至关重要，从而实现计算效率高的操作 [16]。考虑输入 x_k ∈ R^(L×D)，即在信号流中采样长度为 L 的向量（遵循 [17]），根据零阶保持规则（zeroth-order hold rule），可以将 ODE（方程1）离散化如下：

其中，B、C ∈ R^(D×N)，Δ ∈ R^D。在实践中，根据 [12]，我们通过使用一阶泰勒级数来精练对¯B 的近似：

选择性扫描机制。与主要关注线性时不变（linear time-invariant，LTI）SSMs 的普遍方法有所不同，VMamba 通过将选择性扫描机制（S6）[12] 作为核心 SSM 运算符而使自己与众不同。在 S6中，矩阵 B ∈ R^(B×L×N)、C ∈ R^(B×L×N) 和 Δ ∈ R^(B×L×D) 是从输入数据 x ∈ R^(B×L×D) 导出的。这意味着 S6 了解嵌入输入的上下文信息，确保了该机制内部权重的动态性。

3.2 2D 选择性扫描

尽管 S6 具有独特的特性，但它以因果方式处理输入数据，因此只能捕获数据的扫描部分内的信息。这自然地使 S6 与涉及时间数据的 NLP 任务相吻合，但在适应非因果数据（例如图像、图形、集合等）时面临重大挑战。解决这个问题的一种直接方法是沿两个不同方向（即前向和后向）扫描数据，允许它们互相补偿而不增加计算复杂性。

尽管图像具有非因果性质，但与文本不同之处在于它们包含 2D 空间信息（例如局部纹理和全局结构）。为了解决这个问题，S4ND [35] 建议通过卷积对 SSM 进行重新构建，并通过外积将核从 1-D 直接扩展到 2-D。然而，这种修改阻止了权重的动态性（即，与输入无关），导致上下文基础数据建模能力的丧失。因此，我们选择通过坚持选择性扫描方法 [12] 来保留动态权重，这不幸地使我们无法按照 [35] 的方法集成卷积操作。

为了解决这个问题，我们提出了如图 2 所示的交叉扫描模块（Cross-Scan Module，CSM）。我们选择沿行和列展开图像补丁成序列（扫描扩展），然后沿四个不同方向进行扫描：从左上到右下，从右下到左上，从右上到左下，从左下到右上。这样，任何像素（如图 2 中的中心像素）都会从不同方向的所有其他像素中集成信息。然后，我们将每个序列重新整形成单个图像，并将所有序列合并成一个新的序列，如图 3 所示（扫描合并）。

将 S6 与 CSM 集成，称为 S6 块，作为构建视觉状态空间（Visual State Space，VSS）块的核心元素，构成了 VMamba 的基本构建块（在下一小节中进一步详细说明）。我们强调 S6 块继承了选择性扫描机制的线性复杂性，同时保留了全局感受野，这与我们构建这样一个视觉模型的动机一致。

3.3 VMamba 模型

3.3.1 整体架构

VMamba-Tiny 的架构概览如图 4（a）所示。VMamba 从使用一个干节点将输入图像划分为补丁开始，类似于 ViTs，但没有将补丁进一步展平成 1-D 序列。这种修改保留了图像的 2D 结构，导致具有 H / 4 × W / 4 × C_1 维度的特征映射。

然后，VMamba 在特征映射上堆叠多个 VSS 块，保持相同的维度，构成 “Stage 1”。在 VMamba中，通过补丁合并操作 [27] 在 “Stage 1” 中对特征映射进行下采样来构建分层表示。随后，涉及更多的 VSS 块，导致输出分辨率为 H / 8 ×W / 8，形成 “Stage 2”。此过程重复以创建分辨率为 H / 16 × W / 16 和 H / 32 × W / 32的 “Stage 3” 和 “Stage 4”。所有这些阶段共同构建了类似于流行的CNN 模型[19, 22, 41, 29, 42] 和一些 ViTs [27, 48, 6, 56] 的分层表示。生成的架构可以在实际应用中作为对其他视觉模型的多功能替代品。

我们在三个不同的尺度上开发 VMamba，即 VMamba-Tiny、VMamba-Small 和 VMamba- Base（分别称为 VMamba-T、VMamba-S 和 VMamba-B）。详细的架构规格见表 1。所有模型的FLOP 都是使用 224×224 的输入大小评估的。将在未来的更新中引入其他架构，如大规模模型。

3.3.2 VSS块

VSS块的结构如图 4(b) 所示。输入经过初始线性嵌入层，输出分为两个信息流。一个流通过一个 3 × 3 的深度卷积层，然后通过 Silu 激活函数 [37] 进入核心 SS2D 模块。SS2D 的输出通过一层标准化层，然后加到另一个信息流的输出上，该信息流经过 Silu 激活。这种组合产生了 VSS 块的最终输出。

与视觉 transformer 不同，由于 VMamba 的因果性质，我们不使用位置嵌入偏差。我们的设计与典型的视觉 transformer 结构不同，后者在一个块中采用以下操作顺序：Norm→attention→Norm→MLP，并丢弃 MLP 操作。因此，VSS 块比 ViT 块更浅，这使我们能够在相似的总模型深度预算下堆叠更多的块。

4. 实验