【论文阅读】Vision Mamba:双向状态空间模型的的高效视觉表示学习

文章目录

• Vision Mamba:双向状态空间模型的的高效视觉表示学习
• • 摘要
  • 介绍
  • 相关工作
  • • 用于视觉应用的状态空间模型
  • 方法
  • • 准备
    • 视觉Mamba
    • Vim块
    • 结构细节
    • 高效分析
    • • 计算效率
  • 实验
  • • 图片分类
    • 语义分割
    • 目标检测和实例分割
    • 消融实验
    • • 双向SSM
      • 分类设计
  • 总结和未来工作

论文地址：

Vision Mamba

Vision Mamba:双向状态空间模型的的高效视觉表示学习

摘要

在这篇文章，我们展示对于视觉表示学习，不再依赖自注意力，提出一个新的泛用的视觉backbone基于双向Mamba块，通过位置嵌入来标示图片序列，通过双向状态空间模型压缩视觉表示。结果表示，Vim有能力解决在高分辨率图片上应用tranformer类型所导致计算和存储限制，有潜力成为下一代视觉基础模型的backbone。

介绍

最近的工作，Mamba引入时变参数到SSMM而且提出一个硬件感知算法以保持其高效训练和推断，Mamba优秀的缩放表现表明在语言建模是Transformer的有前景的替代。
然而，纯基于SSM的主干网络还没有被探索用于处理视觉数据，例如图片和视频。

相比于卷积神经网络，ViT的核心优势是ViT可以为每张图片切片提供数据\切片依赖的全局上下文信息通过自注意力。不同于卷积核对于所有的位置应用相同的参数。另一个优势是模态不可知建模通过将一张图片当作一系列切片而没有引入2D归纳偏置，使其更适应多模态应用。在同事，Transformer的自注意力机制导致了速度和内存使用的挑战，在处理长距离视觉依赖，例如处理高分辨率图片时。

然而，Mamba有两个主要的挑战，即单向建模和缺乏位置感知。为了解决这些挑战，我们提出了视觉Mamba模型，它结合了用于数据相关的全局视觉上下文建模的双向SSM和用于位置感知视觉识别的位置嵌入。我们首先将输入图片切分为块然后线性映射它们为向量输入Vim。图片块在Vim块中作为序列信息被处理，通过提出的双向选择状态空间模型高效压缩视觉表示。更进一步，Vim块的位置嵌入提供了空间信息的感知，使得Vim对于密集预测任务更加鲁棒。在当前阶段，我们使用ImageNet数据集在有监督图片分类任务上训练Vim模型，之后将预训练Vim作为backbone，为下游密集预测任务即语义分割，目标检测和实例分割实现序列视觉表示学习。如同Transformer，Vim可以在大规模无监督视觉数据集上预训练以获得更好的视觉表示。得益于Mamba更好的效率，Vim可以以更低的计算代价实现。

相比于其他为视觉任务设计的基于SSM模型，Vim是一个纯SSM方法并且以序列风格建模，对于通用高效的backbone来说更有前景。

我们主要的贡献总结如下：

• 我们提出了视觉Mamba（Vim）结合数据依赖的全局上下文建模的双向SSM和位置感知视觉理解的位置嵌入结合。
• 不需要attention。
• 我们在ImageNet分类和密集预测下游任务上完成了大量的实验，展示相比于大家接受和高度优化的plain视觉Transformer例如DeiT的优越性。

相关工作

用于视觉应用的状态空间模型

Long movie clip classification with state-space video models.

使用1D S4处理视频分类的长距离时序依赖

S4nd: ´ Modeling images and videos as multidimensional signals with state spaces.

更近一步拓展1D S4来处理多维数据包括2D图片和3D视频。

Efficient movie scene detection using state-space transformers

结合了S4和自注意力的优势建造了TranS4mer模型，在视频场景检测取得SOTA

Selective structured state-spaces for long-form video understanding

为S4提出了一种新的选择机制，大大提升了S4在长格式视频理解中的表现

Diffusion models without attention

抛弃了注意力机制，使用一个更可缩放的基于SSM的backbone来生成高分辨率图像和在可负担的计算条件下处理细粒度表达。

U-mamba: Enhancing long-range dependency for biomedical image segmentation.

提出了U-Mamba，一个混合CNN-SSM结构，在生物医学图像上处理长距离依赖。

上述的工作要么是应用SSM到特定的视觉应用或者建立一个结合SSM和卷积或注意力的混合结构。和它们不同，我们建立了一个纯-SSM模型，可以作为一个通用视觉backbone

方法

准备

基于SSM的模型即状态空间序列模型（S4）和Mamba受连续系统启发。这一部分介绍和离散化参见

【论文翻译】Mamba 中的状态空间模型背景-CSDN博客

【论文翻译】MODELING SEQUENCES WITH STRUCTURED STATE SPACES 附录A.1-CSDN博客

视觉Mamba

图二展示了Vim的概览，标准的Mamba是为1-D序列设计的。为了处理视觉任务，我们首先将2-D图像到展平的2-D切块，其中是输入图像的尺寸，是通道的数量，是图片切片的尺寸。接着，我们将象形映射为一个长度D的向量

并加入位置嵌入如下：

其中是的第个切片，是可学习的映射矩阵。受ViT和BERT启发，我们也是用类别token表示整个patch序列，表示为。我们之后将token序列送入Vim编码器第层，然后得到输出。最后，我们归一化输出类token并将其送入MLP得到最后的预测如下：

Vim是提出的视觉mamba块，是层的数目，是归一化层

Vim块

原始的Mamba块是为1-D序列设计的，并不适宜用于需要空间感知理的解视觉任务。在这个部分，我们提出里Vim块，包含为视觉任务设计的双向序列建模。Vim块如图二

特别的，我们在算法31展示了Vim块的操作。输入token序列首先被归一化层归一化。接着我们以E维度线性投影归一化序列到x和z。之后我们从前向和后向方向处理x。对于每个方向，我们首先在x应用1-D卷积得到。我们接着线性投影到，被用于变换。最后我们通过SSM计算和之后被变量z门控，并相加得到输出的token序列

结构细节

总结，我们结构的超参数如下：

L：块的数目

D：隐藏状态的维度

E：扩展的状态维度

N：SSM维度

追随ViT和DeiT，我们首先应用一个16×16的映射层得到一个无重叠的切块嵌入的1-D序列。之后，我们直接堆叠个Vim块。在默认情况下，我们将设置为24，SSM维度为16。来对齐DeiT系列的模型大小，对于tiny变种我们将隐藏状态模型D设置为192并且将扩展状态维度设置为384。对于small变种我们将隐藏状态模型D设置为384并且将扩展状态维度​设置为768.

高效分析

传统的基于SSM的方法应用了快速傅里叶变换来加快等式4所示的卷积操作。遂于数据依赖的方法，例如Mamba,SSM操作不再等价于卷积。为了解决这个问题，Mamba和我们提出的Vim选择一个现代硬件友好的方式确保高效。这个优化的关键想法在于避免IO和存储限制。

计算效率

Vim块中的SSM和Transformer中的自注意力都在自适应提供全局上下文中扮演关键角色。给定一个视觉序列和默认设置，自注意力和SSM的计算复杂度为：

实验

图片分类

语义分割

目标检测和实例分割

得益于SSM相对于序列长度的线性复杂度，在输入分辨率提高的情况下，其相对于平方复杂度的Transformer类模型会展现出极大优势。

消融实验

双向SSM

为了全面评估在ImageNet上学习到的表达能力，我们使用一个简单的只有两层的Segmenter头来实现在语义分割上的迁移学习。我们研究了这些双向策略：

• None。我们直接采用Mamba block的前向方向来处理视觉序列
• 双向序列。在训练过程中，我们随机翻转视觉序列，类似于数据增强。
• 双向块。我们将堆叠块配对。每对的第一个块以前向方向处理视觉视觉序列，第二个块以后向方向处理视觉序列
• 双向SSM，对于每个块我们添加了额外的SSM来以后向方向处理视觉序列
• 双向SSM+Conv1d。基于双向SSM，我们进一步在后向SSM前增加了后向Conv1d

如表4所示，直接应用Mamba块在分类中取得了良好的表现。然而，其不自然的非方向性的方式为下游密集预测任务造成了挑战。特别是，原始的使用双向块的双向策略在分类上低了2.3个点，但相比于原始的无方向性Mamba块在语义分割上提升了1.3mIoU。通过加入二外的后向SSM和Conv1d，我们实现了更好的分类准确性和巨大的分类提升。我们使用双向SSM+Conv1d作为Vim块的默认设置。

分类设计

我们在Vim的分类设计上进行了消融实验。我们研究论下列分类策略：

• 平均池化
• 最大池化 我们首先在视觉序列中的每个token上调整分类头，接着在需略上应用最大池化得到预测结果
• 头类别token 跟随DeiT，我们拼接视觉序列地头部的类别token，并进行分类
• 双类别tokrn 基于头部类别token策略，我们在视觉序列的尾部额外添加了一个类别token
• 中类别token 我们在视觉序列的中部添加了一个视觉序列之后再最后的中部类别token上进行分类。

如表格五所示，实验表明，中部类别token策略可以完整的利用SSM天然的循环性和ImageNet的中心物体先验，展示了最好的top-1准确度76.1。

总结和未来工作

我们提出里视觉Mamba（Vim）来探索最近高效的状态空间模型，即Mamba，作为一个通用的视觉backbone。不像先前为视觉任务设计的的状态空间模型使用了混合结构或者等价于全局2D卷积核，Vim实现了数据依赖的全局视觉上下文并且拥有和Transformer一般的建模能力，但拥有更小的计算复杂度。受益于Mamba的硬件感知设计，推理速度和内存使用显著优于ViT在处理高分辨率的图片情况下。实验结果证明了Vim的建模能力和高效，展示了Vim在成为下一代视觉backbone方面的巨大潜力。

在未来的工作，位置编码的双向SSM建模Vim对于例如掩膜图像建模预训练和由相似结构的Mamba可以适合多模态任务例如CLIP类型的预训练。基于预训练的Vim权重，探索在高分辨率医学图像，遥感图像和长视频方面，即下游任务Vim的用处，是显而易见的。