视觉 Mamba 的综述

24年4月来自香港科技大学的论文“Visual Mamba: A Survey and New Outlooks”。

Mamba 是一种近期出现的选择性结构化状态空间模型，它在长序列建模方面表现出色，这在大模型时代至关重要。长序列建模带来了重大挑战，包括捕获数据中的长距离依赖关系以及处理由其过长长度引起的计算需求。Mamba 通过克服卷积神经网络的局部感知限制和 Transformers 的二次项计算复杂度来解决这些挑战。鉴于其优于这些主流基础架构的优势，Mamba 展现出成为视觉基础架构的巨大潜力。Mamba 已积极应用于各种计算机视觉任务，并产生了许多贡献。

本文回顾视觉 Mamba 方法，分析 200 多篇论文。概述原始 Mamba 模型的公式。深入研究代表性主干网络和用不同模态分类其应用程序，包括图像、视频、点云和多模态。特别是，扫描技术对于 Mamba 适应视觉任务至关重要，将这些扫描技术解偶，明确其功能并增强其在各种应用程序中的灵活性。最后，讨论挑战和未来方向。

由于状态表示需要大量的计算和内存，SSM 在实践中并未得到广泛应用。随着结构化 SSM (S4) 的出现，这种情况发生了改变，它通过重新参数化状态矩阵 (Gu et al, 2022b) 解决了这些限制。此后，出现了一系列包含 SSM 层的 SSM 变换和神经网络架构 (Gu et al, 2022b,a; Gupta et al, 2022; Orvieto et al, 2023; Smith et al, 2023; Gu et al, 2023; Hasani et al, 2023; Fu et al, 2023)。然而，SSM 的恒定序列转换限制其基于上下文的推理能力，这一点在有效的Transformer 等模型至关重要 (Vaswani et al, 2017)。

Mamba（Gu & Dao，2023）将选择机制集成到 SSM 中来解决此问题，从而使 SSM 能够根据当前 token 沿序列或扫描路径选择性地传播或忘记信息。此外，为了有效地计算这些选择性 SSM，开发一种硬件-觉察算法。将这些选择性 SSM 集成到简化的神经网络架构中，称为 Mamba。凭借类似于 Transformers 的建模能力和序列长度的线性可扩展性，Mamba 成为序列建模有前途的基础架构，这在大模型时代至关重要。

由于序列建模或自然语言处理 (NLP) 技术越来越多地应用于 CV，Mamba (Gu & Dao, 2023) 迅速应用于 CV 任务 (Zhang et al, 2024a; Liu et al, 2024f; Heidari et al, 2024; Zou et al, 2024b)。VMamba (Liu et al, 2024g) 是一种早期的代表性视觉 Mamba 模型。它沿图像的水平和垂直维度将图像块展开为序列，并沿这两个方向进行双向扫描。另一个视觉 Mamba 模型 Vim (Zhu et al, 2024a) 受 ViT (Dosovitskiy et al, 2021) 的启发，利用位置嵌入来整合空间信息。它还使用双向 SSM 来处理非因果图像序列。同样，其他一些值得注意的研究（Li et al, 2024d; Yang et al, 2024a; Patro & Agneeswaran, 2024b; Huang et al, 2024e; Pei et al, 2024）也深入探索了视觉主干网络，在分类、检测和分割任务中始终如一地实现了有竞争力的性能。除了这些努力之外，Mamba 还应用于各种视觉模态及其各自的应用，涵盖图像处理、视频分析、点云处理和多模态场景。凭借上述建模能力和线性可扩展性，Mamba 成为 CV 任务有前途的基础架构。

状态空间模型 SSM

S4 (Gu et al, 2022b) 中的 SSM 变换，源自经典的状态空间模型，该模型将一维输入信号 x(t) 通过 潜状态 h(t) 映射到一维输出信号 y(t)。该变换在数学上可表示为线性常微分方程 (ODE)：

为了处理离散输入序列 x = (x0,x1,...)，可以使用各种离散化规则以步长 ∆ 离散化公式 (1) 中的参数，该步长可以看作是连续输入 x(t) 的分辨率。根据先前的研究 (Tustin, 1947)，S4 (Gu et al, 2022b) 使用双线性方法离散化这些参数。根据 (Gupta et al, 2022)，Mamba (Gu & Dao, 2023) 采用零阶保持 (ZOH) 假设来求解 ODE。具体而言，这会导致连续参数 A、B 转换为其离散对应项 A、B，如下所示：

这样离散化方程（1）为如下：

等式 (3) 表示从 xt 到 yt 的序列到序列映射。此配置允许将离散化 SSM 计算为 RNN。但是，由于其顺序性，这种离散化循环 SSM 不适用于训练。

为了实现高效的并行训练，这个递归过程也可以重新表述并计算为一次卷积（Gu et al, 2022b）如下：

选择性 SSM

SSM 由公式 (1)、公式 (3) 或公式 (4) 表示的参数相对于输入或时间动态保持不变。Mamba (Gu & Dao, 2023) 将此 线性时不变 (LTI) 属性视为 SSM 在基于上下文推理方面的一个基本限制。为了解决这个问题，Mamba 采用了选择机制。选择机制的实现方法是简单地将 SSM 的参数配置为输入的函数，从而实现沿序列的输入-相关交互。具体而言，参数 B、C、∆ 取决于输入序列 x：

其中 sB (x) = LinearN (x) 和 sC (x) = LinearN (x) 都将输入投影到一个维度 N。s∆(x) = BroadcastD(Linear1(x)) 首先将输入投影到维度 1，然后将其广播到维度 D。τ∆ 是 softplus 函数。x 大小是B×L×D，其中 B 表示批量大小，L 是序列长度，D 是通道数。

由此产生的选择性 SSM 无法通过 RNN 或卷积来计算。Mamba (Gu & Dao，2023) 采用硬件-觉察算法高效计算选择性 SSM。硬件-觉察算法利用三种经典技术：并行扫描、核融合和重新计算。

Mamba

Mamba 是一种简化的 SSM 架构，借鉴了 H3 架构（Fu et al, 2023）的灵感，后者是许多 SSM 架构的基础。具体而言，H3（Fu et al, 2023）块集成两个离散的 SSM 来明确模拟线性注意。常用的 SSM 架构交替堆叠 H3 块（Fu et al, 2023）和流行的 MLP 块。与它们不同，Mamba 将这两个块集成在一起，构建了简化的 Mamba 块。如图所示，可以从两个不同的角度来查看 Mamba 块。首先，它用激活函数替换线性注意类或 H3（Fu et al, 2023）块中的第一个乘法门。其次，它将 SSM 变换合并到 MLP 块的主要通路中。Mamba 的整体架构由重复的 Mamba 块与标准规范化层和残差连接交错组成。

Mamba 继承了状态空间模型在序列长度上的线性可扩展性，同时也实现 Transformers 的建模能力。Mamba 展现了 CV 中两种主要基础架构（即 CNN 和 Transformers）的显著优势，使其成为 CV 一个有前途的基础架构。与依赖于显式存储整个上下文进行基于上下文推理的 Transformers 不同，Mamba 采用一种选择机制，该机制根据当前 token 有选择地沿序列传播或过滤信息。这确保了建模过程仅受序列中过去和当前输入的影响，从而遵循因果关系原则。因此，这种选择机制的1D和因果特性，成为将 Mamba 应用于 CV 的研究人员的关注重点。

纯Mamba

Vim

Vim (Zhu et al, 2024a) 是一种基于 Mamba 的架构，直接对类似于 ViT (Dosovitskiy et al, 2021) 的图像块序列进行操作。首先将输入图像转换为扁平的 2D 块，然后使用线性投影层对其进行矢量化，并添加位置嵌入以保留空间信息。继 ViT (Dosovitskiy et al, 2021) 和 BERT (Devlin et al, 2019) 之后，将一个类token附加到patch token序列中。然后将整体 token 序列输入 Vim 编码器，该编码器的结构为一系列相同的 Vim 块。如图所示，Vim 块是一个 Mamba 块，它将后向 SSM 路径与前向 SSM 路径集成在一起。

VMamba (Liu et al, 2024g) 发现，由于其选择机制具有 1D 因果属性，在将 Mamba 应用于 2D 图像时，方向敏感的挑战是存在的。选择机制使 Mamba 无法吸收未扫描数据部分的信息。因此，1D 扫描对于同时捕获图像不同方向的依赖信息并不是最佳选择，从而导致感受野受限。

为了应对这一挑战，VMamba 引入了 交叉-扫描模块 (CSM) 。CSM 将输入图像转换为沿水平和垂直轴的patches序列（交叉-扫描），从四个方向扫描序列：从左上到右下、从右下到左上、从右上到左下、从左下到右上。然后使用选择性 SSM 分别处理得到的四个序列。此操作也可以看作是沿 2D 轴执行双向选择性 SSM。因此，每个像素在四个不同方向上整合来自所有其他像素的信息。最后，所有序列都转换为其原始 2D 布局以构成地图并合并获得输出地图（交叉-合并）。交叉-扫描、选择性 SSM 和交叉-合并的整个过程，构成了 2D- 选择性-扫描 (SS2D)。

VMamba 首先将输入图像转换为 2D patches。然后将它们输入到 VMamba 的多个阶段中，这些阶段由堆叠的 视觉状态空间 (VSS) 块组成，并在这些阶段之间插入下采样操作 (Liu et al, 2022a) 以构建整体分层架构。VMamba 引入了两种 VSS 块，即 vanilla VSS 块和 VSS 块。原始 VSS 块与 Mamba 块类似，但用 2D 深度卷积层代替 1D 卷积层，用 SS2D 和层归一化 (LN) 层代替了选择性 SSM。VSS 块类似于典型的 Transformer 块。

Plain Mamba (Yang et al, 2024a) 被设计为一种非层次结构，实现以下几个目标：(1) 非层次结构有利于多层特征融合，增强不同尺度的集成；(2) 它支持多模态数据的有效融合；(3) 其更简单的架构往往提供更好的泛化能力；(4) 它适合硬件加速优化。

由于视觉数据缺乏预定义的顺序及其固有的空间维度，一些纯 Mamba 网络专注于扫描技术。Mamba-ND (Li et al, 2024d) 旨在将 Mamba 扩展到包括图像和视频在内的多维数据。它将一维 Mamba 层视为黑盒子，并探索如何解开和排序多维数据。结果表明，一组Mamba 层和简单的交替-方向性排序可实现卓越的性能。FractalMamba (Tang et al, 2024b) 采用分形扫描曲线，即 希尔伯特曲线 来处理二维图像块。不同的是，Mamba® (Wang et al, 2024a) 发现在视觉 Transformers 中观察的特征伪影，在 Vim (Zhu et al, 2024a) 中更为明显。为了缓解这个问题，Mamba® 通过策略性地将输入无关的tokens，称为寄存器（Darcet，2024）tokens，插入到视觉tokens序列中来改进 Vim。ARM（Ren，2024b）表明，自回归预训练非常适合 Mamba 架构，可显著增强其视觉性能并释放其扩展潜力。

Hybrid Mamba

除了独立使用外，一些研究还探索了将 Mamba 与其他架构范例集成以构建混合 Mamba 网络。

Local Mamba

Local Mamba (Huang et al, 2024e) 确定了 Vim (Zhu et al, 2024a) 和 VMamba (Liu et al, 2024g) 模型中的一个重大限制，即在单次扫描过程中空间局部tokens之间的依赖关系被破坏。为了解决这个问题，如上图所示的局部采样，Local Mamba 将输入图像分成多个局部窗口，像 VMamba (Liu et al, 2024g) 一样在不同方向上执行 SSM，同时还保持全局 SSM 操作。此外，Local Mamba 在patch合并之前实现了空间和通道注意模块，增强方向性特征的集成并减少冗余。Local Mamba 采用了一种策略来为每一层选择最有效的扫描方向，从而优化计算效率。

EfficientVMamba (Pei et al, 2024) 引入高效 2D 扫描 (ES2D) 技术，该技术采用特征图上的patch空洞采样来减少计算负担。ES2D 用于提取全局特征，而并行卷积分支用于提取局部特征。然后，两种特征类型都由收缩和激励 (SE，Squeeze-and- Excitation) 块单独处理。总的来说，ES2D、卷积分支和 SE 块构成高效视觉状态空间 (EVSS) 块的核心组件。EVSS 块的输出，是调制后全局特征和局部特征的总和。EVSS 块构成了 EfficientVMamba 的早期阶段，而 Inverted Residual 块（Sandler，2018）则反向构成了后期阶段。