论文阅读：VMamba: Visual State Space Model

论文阅读：VMamba: Visual State Space Model

abstract

卷积神经网络（CNNs）和ViTs是视觉表示学习的两个最受欢迎的基础模型。尽管CNNs在图像分辨率方面表现出显著的可扩展性，具有线性复杂度，但ViTs在拟合能力上超越了它们，尽管它们面临着二次方复杂度的挑战。仔细检查发现，ViTs通过整合全局感受野和动态权重实现了卓越的视觉建模性能。这一观察启发我们提出了一种新的架构，它继承了这些组件，同时提高了计算效率。为此，我们从最近引入的状态空间模型中汲取灵感，提出了视觉状态空间模型（VMamba），它在不牺牲全局感受野的情况下实现了线性复杂度。为了解决遇到的与方向敏感性相关的问题，我们引入了交叉扫描模块（CSM）来遍历空间域，并将任何非因果视觉图像转换为有序的块序列。

3. Method

3.1 Preliminaries

State Space Models

State Space Models (SSMs) 通常被视为线性时不变系统，它们将激励$x(t)\in {\mathbb{R}}^L$映射到响应$y(t)\in {\mathbb{R}}^L$。数学上，这些模型通常被表述为线性常微分方程（ODEs），公式如下：

$$\begin{array}{rl}{h}^{\prime }(t)& =Ah(t)+Bx(t)\\ y(t)& =Ch(t)+Dx(t)\end{array}$$

其中，参数包括状态大小为$N$的矩阵 $A\in {\mathbb{C}}^{N\times N},B,C\in {\mathbb{C}}^N$以及跳跃连接 $D\in {\mathbb{C}}^1$。

Discretization

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

离散化的主要目标是将ODE转换为离散函数。这种转换对于将模型与输入数据中体现的信号的采样率对齐至关重要，从而实现计算效率的操作。考虑到输入$x_k\in {\mathbb{R}}^{L\times D}$，一个在长度为 L 的信号流中的采样向量，ODE可以使用零阶保持规则进行如下离散化：
$$\begin{array}{rl}& h_k=\bar{A}{h}_{k-1}+\bar{B}{x}_k,\\ & \begin{array}{r}{y}_k=\bar{C}{h}_k+\bar{D}{x}_k,\end{array}\\ & \bar{A}=e^{\Delta A},\\ & \bar{B}=(e^{\Delta A}-I)A^{-1}B,\\ & \bar{C}=C\end{array}$$
其中，$B,C\in {\mathbb{R}}^{D\times N}$，并且$\Delta \in {\mathbb{R}}^D$。在实践中，遵循文献[12]，我们使用一阶泰勒级数改进$\bar{B}$的近似：
$$\bar{B}=(e^{\Delta A}-I)A^{-1}B\approx (\Delta A)(\Delta A{)}^{-1}\Delta B=\Delta B$$

Selective Scan Mechanism

与主要关注线性时不变（LTI）SSMs的流行方法不同，VMamba通过将选择性扫描机制（S6）作为核心SSM操作符。在S6中，矩阵$B\in {\mathbb{R}}^{B\times L\times N},C\in {\mathbb{R}}^{B\times L\times N}$ 和 $\Delta \in {\mathbb{R}}^{B\times L\times D}$是从输入数据$\begin{array}{r}x\in {\mathbb{R}}^{B\times L\times D}\end{array}$派生出来的。这意味着S6能够意识到嵌入在输入中的上下文信息，确保了该机制中权重的动态性。

3.2 2D Selective Scan

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

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

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

信息流的比较：注意模块与交叉扫描模块(CSM)(A)注意机制对中心象素的所有象素进行统一积分，导致O(N 2)复杂度。(B)CSM集成了左上角、右下角、右上角和左下角的像素，复杂度为O(N)。

S6与CSM的整合，被称为S6块，是构建视觉状态空间（VSS）块的核心元素，而VSS块构成了VMamba的基本构建模块（在下一子节中有更详细的介绍）。我们强调，S6块继承了选择性扫描机制的线性复杂度，同时保留了全局感受野，这与我们构建此类视觉模型的动机相一致。

3.3 VMamba Model

3.3.1 Overall Architecture

VMamba-Tiny的架构概述如图所示。VMamba首先通过类似于ViTs的stem模块将输入图像分割成小块，但不进一步将这些小块展平成1-D序列。这种修改保留了图像的2-D结构，结果是一个具有$\frac{H}{4}\times \frac{W}{4}\times C_1$维度的特征图。

VMamba然后在特征图上堆叠多个VSS块，保持相同的维度，构成“第1阶段”。通过在“第1阶段”中进行小块合并操作[27]，构建了VMamba的层次化表示。随后，涉及更多的VSS块，输出分辨率为$\frac{H}{8}\times \frac{W}{8}$，形成“第2阶段”。这个过程重复进行，创建出分辨率为$\frac{H}{16}\times \frac{W}{16}$的“第3阶段”和$\frac{H}{32}\times \frac{W}{32}$的“第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中概述。所有模型的FLOPs都是使用224×224输入尺寸进行评估的。未来更新将引入其他架构，如大规模模型。

3.3.2 VSS Block

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

与ViTs不同，我们没有在VMamba中使用位置嵌入偏置，因为它具有因果性质。我们的设计与传统的ViT结构不同，后者在块中采用以下操作序列：Norm → 注意力 → Norm → MLP，并丢弃了MLP操作。因此，VSS块比ViT块更浅，这使我们能够在保持总模型深度预算相似的情况下堆叠更多的块。