Vision Mamba 论文学习精读笔记 Efficient Visual Representation Learning with Bidirectional State Space Model

论文地址：Vision Mamba
代码地址：Vim

摘要

总结了论文中提出的Vision Mamba（Vim）模型的动机、设计以及其在视觉任务上的表现。

首先，指出了状态空间模型（SSM）在长序列建模方面的潜力，尤其是Mamba模型在硬件效率方面的优势。纯粹基于 SSM 构建高效且通用的视觉 backbone 是一个有吸引力的方向。

然后提到，尽管在视觉数据表示上存在挑战，例如视觉数据的位置敏感性和视觉理解对全局上下文的需求，但作者成功地展示了不必完全依赖自注意力机制来学习视觉表示。

文章介绍了Vim模型，这是一种新的通用视觉基础模型，它利用双向Mamba块(bidirectional Mamba blocks (Vim))和位置嵌入 (position embeddings)来处理图像序列，并在ImageNet分类、COCO对象检测和ADE20K语义分割任务上取得了比现有的视觉Transformer模型（如DeiT）更好的性能。

Vim在处理高分辨率图像时，不仅计算和内存效率显著提高（例如，比DeiT快2.8倍，GPU内存减少86.8%），而且还能够克服Transformer在高分辨率图像理解上的计算和内存限制，显示出成为下一代视觉基础模型的巨大潜力。

1. 引言

状态空间模型（SSM）的发展和在长序列数据建模中的应用，以及它在计算机视觉领域的潜力和挑战

• SSM的发展: SSM的概念起源于经典的卡尔曼滤波器，这是一种用于估计线性动态系统状态的算法。近年来，SSM在捕捉长距离依赖性方面显示出优越性，并且能够支持并行训练。
• SSM的应用: 文中提到了几种基于SSM的方法，如线性状态空间层（LSSL）、结构化状态空间序列模型（S4）、对角状态空间（DSS）和S4D，它们在处理具有长距离依赖性的序列数据方面表现出色。
• SSM与其他架构的结合: SSM也与其他架构，如卷积神经网络（CNN）和Transformer结合，用于处理二维数据。例如，2-D SSM、SGConvNeXt和ConvSSM结合了SSM和CNN或Transformer结构。
• Mamba的创新: Mamba是近期的一个创新，它引入了时间变化参数到SSM，并提出了一个硬件感知的算法，使得训练和推断变得非常高效。Mamba在扩展性能上的优越性表明它是语言建模中Transformer的一个有前景的替代品。
• 视觉数据处理的挑战: 尽管SSM在处理序列数据方面表现出色，但它在处理视觉数据（如图像和视频）方面还没有一个通用的、纯粹基于SSM的后端网络。这暗示了在视觉领域开发SSM方法还有很大的潜力和挑战。

简而言之，强调了SSM在捕获长范围依赖方面的优势，Mamba模型在此基础上的进一步发展，以及将SSM应用于视觉数据处理的未来潜力。

视觉Transformer（ViT）在视觉表示学习中的成功及其优势和挑战

• ViT的成功：ViT在视觉表示学习方面取得了巨大成功，特别是在大规模自监督预训练和下游任务的高性能表现上。
• ViT的核心优势：
  • 数据/块依赖的全局上下文：ViT能够通过自注意力机制为每个图像块 (patch) 提供依赖于数据或块的全局上下文。这与卷积神经网络（CNN）不同，CNN在所有位置使用相同的参数（即卷积滤波器）。
  • 模态不可知建模：ViT将图像视为一系列块 (patches)，不包含 2D 归纳偏置，这种处理方式使得 ViT 成为多模态应用的首选架构，因为它可以更灵活地处理不同类型的数据。
• ViT的挑战：
  • 尽管ViT在提供全局上下文方面非常有效，但其自注意力机制在处理具有长距离视觉依赖的高分辨率图像时，会面临速度和内存使用方面的挑战。这是因为自注意力的计算复杂度与输入长度的平方成正比，导致在处理大型图像或长序列数据时效率降低。

总的来说，强调了ViT在全局上下文理解和多模态应用中的优势，同时也指出了它在处理复杂视觉任务时可能遇到的性能瓶颈。

如何将Mamba模型从语言建模领域转移到视觉领域，以及Vision Mamba（Vim）模型的设计理念、挑战和应用

这段话讨论了如何将Mamba模型从语言建模领域转移到视觉领域，并介绍了Vision Mamba（Vim）模型的设计理念、挑战和应用。

• 从语言到视觉的转移：由于Mamba在语言建模中取得的成功，作者受到启发，希望将这种成功转移到视觉领域，设计一种通用而高效的视觉基础架构 (backbone)，使用先进的状态空间模型（SSM）方法。
• Mamba的挑战：
  • 单向建模：Mamba原本用于语言处理，通常是单向的，这意味着它只能捕捉从前到后的序列依赖。
  • 位置感知的缺失：Mamba在处理视觉数据时，需要对空间位置有所感知，这是它原本的设计中缺乏的。
• Vision Mamba (Vim)模型的提出：为了解决这些挑战，提出了Vim模型。Vim引入了双向SSM来进行数据依赖的全局视觉上下文建模，并加入位置嵌入来实现位置感知的视觉识别。
• Vim的实现细节：
  • 将输入图像切分成块 (patches)，然后线性投影成向量，作为Vim的处理单元。
  • 在Vim块中，图像块被视为序列数据，通过双向选择性状态空间高效地压缩视觉表示。
  • Vim块中的位置嵌入增强了模型对空间信息的感知能力，使其在密集预测任务中更为鲁棒。
• Vim的训练与应用：
  • Vim首先在ImageNet数据集上进行监督图像分类任务的训练。
  • 预训练的Vim模型随后用作各种下游密集预测任务（如语义分割、对象检测和实例分割）的 backbone。
  • 类似于Transformer模型，Vim也可以在大规模无监督视觉数据上进行预训练，以实现更好的视觉表示。
  • 由于Mamba的高效性，Vim能够以较低的计算成本实现大规模预训练。

简而言之，解释了如何将Mamba的成功应用于视觉领域，提出了解决其在视觉任务中存在的挑战的Vim模型，并描述了Vim的实现和训练过程，以及其在各种视觉任务中的应用。

Vision Mamba（Vim）模型在处理视觉任务方面相比其他基于状态空间模型（SSM）和Transformer模型的优势

Vision Mamba（Vim）模型在处理视觉任务方面相比其他基于状态空间模型（SSM）和Transformer模型的优势，尤其是在图像分类、处理高分辨率图像以及执行密集预测任务方面。

• Vim与其他SSM模型的比较：Vim是一种纯粹基于SSM的方法，以序列方式建模图像，这使得Vim成为一个通用且高效的 backbone。与其他基于SSM的视觉模型相比，Vim的这一特点更有利于处理各种视觉任务。
• Vim的独特之处：得益于双向压缩建模和位置感知能力，Vim成为第一个能够处理密集预测任务的纯SSM模型。这表明Vim能够理解和处理图像中的空间关系，为每个图像块提供全局上下文信息。
• Vim与DeiT的比较：与目前最令人信服的基于Transformer的模型（例如DeiT）相比，Vim在ImageNet图像分类任务上表现更优。此外，在处理高分辨率图像方面，Vim在GPU内存使用和推断时间上都更为高效。
• 高效的序列视觉表示学习：Vim的内存和速度效率使其能够直接执行序列视觉表示学习，无需依赖于2D先验知识（例如ViTDet中的2D局部窗口），同时在高分辨率视觉理解任务上比DeiT实现更高的准确率。

简言之，Vim通过其独特的双向SSM结构和对空间信息的有效编码，在视觉表示学习的通用性、效率和性能方面超越了现有的SSM和Transformer模型，尤其是在处理复杂视觉任务和高分辨率图像时的表现。

Vision Mamba（Vim）论文的主要贡献

• Vim的提出：首先，作者提出了Vision Mamba（Vim），这是一种集成了双向状态空间模型（SSM）的新模型，用于数据依赖的全局视觉上下文建模，并加入位置嵌入来实现对空间信息敏感的视觉理解。这表明Vim能够捕获和理解图像中的全局和局部信息，对位置进行编码，从而提高模型对视觉数据的理解能力。

• Vim的计算和内存效率：Vim的一个显著特点是它不依赖于传统的注意力机制，但仍保持与视觉Transformer（ViT）相当的建模能力，同时具有低于二次时间的计算复杂度和线性内存复杂度。具体来说，在处理1248×1248分辨率的图像进行批量推断时，Vim的速度比DeiT快2.8倍，并且节省了86.8%的GPU内存。这一点对于需要处理大量高分辨率图像的应用尤为重要，表明Vim能够在保持高性能的同时大幅降低计算资源消耗。

• 广泛的实验验证：作者在ImageNet分类任务和多个密集预测下游任务上对Vim进行了广泛测试。实验结果显示，Vim在性能上超过了现有的高度优化的视觉Transformer模型，即DeiT。这证明了Vim不仅在理论上具有创新性，在实际应用中也具有显著的性能优势。

2. 相关工作

CV领域中用于视觉表示学习的不同架构的发展历程

Vision Mamba（Vim）是如何从这些发展中获得灵感并提出一种新的架构的。

• 早期架构：最初，卷积神经网络（ConvNet）是计算机视觉的标准网络设计，为各种视觉应用提供了视觉基础模型。随后，提出了多种卷积神经网络架构，以提高性能和效率。

• Vision Transformer（ViT）的创新：ViT改变了传统视觉模型的格局，通过将图像视为一系列扁平化的2D块并直接应用纯Transformer架构，使得在图像分类任务上取得了令人惊讶的结果，促进了许多后续工作的产生。这些工作包括将2D卷积先验引入ViT的混合架构设计，以及改进传统2D卷积网络的工作。

• Transformer模型的挑战与发展：尽管Transformer模型因其较大的模型容量、统一的多模态表示能力、以及对自监督学习友好等优势而成为首选架构，但其面临的一个主要挑战是处理视觉任务时的计算复杂度问题，这限制了视觉标记的数量。针对这一挑战，已有多种工作尝试寻找解决方案，但很少有工作专注于视觉应用。

• Vim的灵感与目标：受到Mamba模型在处理序列数据方面高效性的启发，本工作探索了构建一种不使用注意力机制的纯SSM（状态空间模型）基础模型作为通用视觉基础架构的可能性。该模型旨在保留ViT的序列化、模态不可知的建模优点，同时解决Transformer在视觉表示学习中的计算复杂度问题。

状态空间模型（SSM）在长序列建模领域的应用及其发展

• S4模型：首先介绍了结构化状态空间序列（S4）模型，这是一种新的替代传统CNN和Transformer的模型，能够模拟长距离依赖性。S4模型的一个重要特性是，其计算复杂度随序列长度线性增长，这一点对于长序列数据处理尤为重要。

• S5层：进一步的探索中，提出了S5层，通过引入多输入多输出（MIMO）SSM和高效的并行扫描技术到S4层，以增强模型的处理能力。

• H3层：另一个重要的发展是H3层的设计，该层几乎填补了SSM和Transformer注意力机制在语言模型中的性能差距，展示了SSM在捕获长距离依赖性方面的潜力。

• 门控状态空间层：通过在S4基础上引入更多的门控单元构建了门控状态空间层，以提高模型的表达能力。

• Mamba模型：最近，提出了依赖于数据的SSM层并构建了Mamba模型，这是一个通用的语言模型 backbone，它在处理大规模实际数据时能够超越各种大小的Transformer模型，并且同样享有随序列长度线性增长的计算复杂度。

在这项工作中，作者探索了将Mamba模型的成功转移到视觉领域的可能性，即构建一个纯粹基于SSM而非依赖于注意力机制的通用视觉基础架构。这表明，借鉴SSM在语言模型中的成功经验，通过适当的改进和调整，有可能为视觉任务提供一个高效且能够处理复杂依赖性的新型架构。

状态空间模型（SSM）在视觉应用中的不同应用和发展，以及如何基于这些进展提出了一个纯粹基于SSM的模型

• 1D S4在视频分类中的应用：提到使用一维S4模型来处理视频分类任务中的长期时间依赖问题。

• 扩展至多维数据：介绍了将1D S4扩展到包括二维图像和三维视频在内的多维数据处理。

• TranS4mer模型：结合了S4模型和自注意力机制的优势，创建了TranS4mer模型，该模型在电影场景检测任务上取得了最先进的性能。

• S4的选择性机制：通过引入一种新颖的选择性机制到S4中，显著提高了S4在长视频理解任务上的性能，并大幅降低了内存占用。

• 更可扩展的SSM基础架构：提出了一种用SSM基础架构替代注意力机制的方法，以生成高分辨率图像并在可承受的计算成本下处理细粒度表示。

• U-Mamba模型：提出了一个混合的CNN-SSM架构，名为U-Mamba，用于处理生物医学图像分割中的长距离依赖问题。

这段讨论指出，虽然之前的工作要么将SSM应用于特定的视觉应用，要么通过结合SSM与卷积或注意力机制构建混合架构，但与之不同的是，作者构建了一个纯粹基于SSM的模型，旨在作为一个通用的视觉 backbone。这表示Vision Mamba（Vim）模型旨在不依赖于传统注意力机制的情况下，通过利用SSM的优势，提供一个能够高效处理图像和视频中长距离依赖关系的强大工具，这在视觉任务中是非常重要的。

3. 方法

3.1 Preliminaries

这段内容介绍了状态空间模型（SSM）及其在S4和Mamba模型中的离散化形式，包括它们如何从连续系统的理论启发而来的。

• 状态空间模型：SSM是一种将一维函数或序列 $x(t)\in \mathbb{R}$ 通过隐藏状态 $h(t)\in {\mathbb{R}}^{\text{N}}$ 映射到 $y(t)\in \mathbb{R}$ 的系统。系统利用演化参数 $\mathbf{A}\in {\mathbb{R}}^{\text{N}\times \text{N}}$ 和投影参数 $\mathbf{B}\in {\mathbb{R}}^{\text{N}\times 1}$和 $\mathbf{C}\in {\mathbb{R}}^{1\times \text{N}}$。

• 连续系统的数学形式：连续系统可以用以下微分方程来表示：

  • $h^{\prime }(t)=\mathbf{A}h(t)+\mathbf{B}x(t)$
  • $y(t)=\mathbf{C}h(t)$
    
    其中，$h^{\prime }(t)$ 表示隐藏状态随时间的变化率，$\mathbf{A}$ 是隐藏状态的演化矩阵，$\mathbf{B}$ 是输入信号 $x(t)$ 对隐藏状态的影响矩阵，$\mathbf{C}$ 是输出矩阵。

• 离散化的状态空间模型：S4和Mamba模型是连续系统的离散版本，引入了一个时间尺度参数 $\Delta$ 来将连续参数 $\mathbf{A}$ 和 $\mathbf{B}$ 转换为离散参数 $\overline{\mathbf{A}}$ 和 $\overline{\mathbf{B}}$。

• 离散化的方法：通常用于变换的方法是零阶保持（Zero-Order Hold, ZOH），定义如下：

  • $\overline{\mathbf{A}}=\exp (\Delta \mathbf{A})$
  • $\overline{\mathbf{B}}=(\Delta \mathbf{A}{)}^{-1}(\exp (\Delta \mathbf{A})-\mathbf{I})\cdot \Delta \mathbf{B}$

  这里，$\exp (\cdot )$ 表示矩阵指数，是将线性系统的连续时间演化转换为离散时间的标准方法。$\Delta A$ 和 $\Delta B$ 是连续演化矩阵和输入影响矩阵的离散等效物。$I$ 是单位矩阵。

这个转换是必要的，因为在实际的数据处理和计算机模拟中，我们通常是在离散时间点上工作，而不是连续时间。这样的离散化使得SSM可以用于数字计算和实时数据处理。

然后对连续时间状态空间模型的离散化版本进行了说明。当 $\mathbf{A}$ 和 $\mathbf{B}$ 被离散化后，原本连续时间的动态方程被转换为离散时间的方程，以适应在数字计算中的使用。这里的 $\overline{\mathbf{A}}$, $\overline{\mathbf{B}}$, 和 $\mathbf{C}$ 是离散化版本的参数。

• 离散化的状态方程：离散时间版本的方程如下所示：
  • $h_t=\overline{\mathbf{A}}{h}_{t-1}+\overline{\mathbf{B}}{x}_t$
  • $y_t=\mathbf{C}{h}_t$

这里：

• $h_t$ 代表在时间步长 $t$ 时的隐藏状态。
• $x_t$ 代表在时间步长 $t$ 时的输入。
• $y_t$ 代表在时间步长 $t$ 时的输出。
• $\overline{\mathbf{A}}$ 是状态转移矩阵，它定义了隐藏状态如何随时间演化。
• $\overline{\mathbf{B}}$ 是控制矩阵，它定义了当前输入如何影响当前的隐藏状态。
• $\mathbf{C}$ 是输出矩阵，它将当前的隐藏状态转换为输出。

离散时间的状态空间模型是为了在数字系统中模拟连续系统的行为。在计算机视觉等数字信号处理应用中，这种模型特别有用，因为它可以捕捉输入数据序列随时间的动态变化。在Vision Mamba论文中，这种离散化的状态空间模型被用作视觉任务的建模工具，特别是处理序列数据，例如视频或者由序列组成的图像。
离散时间的状态空间模型是为了在数字系统中模拟连续系统的行为。在计算机视觉等数字信号处理应用中，这种模型特别有用，因为它可以捕捉输入数据序列随时间的动态变化。在Vision Mamba论文中，这种离散化的状态空间模型被用作视觉任务的建模工具，特别是处理序列数据，例如视频或者由序列组成的图像。

最后，模型通过全局卷积从状态空间模型（SSM）中计算输出：
$$\overline{\mathbf{K}}=(\mathbf{C}\overline{\mathbf{B}},\mathbf{C}\overline{\mathbf{AB}},\dots ,\mathbf{C}{\overline{\mathbf{A}}}^{\mathbf{M}-1}\overline{\mathbf{B}}$$$$\mathbf{y}=\mathbf{x}\ast \overline{\mathbf{K}}$$
具体来说，它定义了一个结构化的卷积核 $\overline{\mathbf{K}}$