论文阅读：Vision Mamba- Efficient Visual Representation Learning with Bidirectional State Space Model

文章介绍

• 本文由华中科技大学、地平线、智源人工智能研究院等机构合作；
• 提出了一种带有双向Mamba块（Vim）的新通用视觉baseline，它用位置嵌入标记图像序列，并用双向状态空间模型压缩视觉表示。

问题引入

• 在处理图像和视频等视觉数据方面,基于纯 SSM 的通用baseline尚未得到探索；
• Visual Transformers（ViT）在视觉任务方面表现出色，但是Transformers 中的自注意力机制在处理长距离视觉相关性(如处理高分辨率图像)时，在速度和内存使用方面表现不佳；
• 想借鉴语言建模上Mamba的优势，但是Mamba 有两个问题，即单一直观建模和缺乏位置感知。

设计思路

• 提出一种模型，这种模型结合了双向 SSM，用于数据相关的全局视觉上下文建模，以及用于位置感知视觉识别的位置嵌入；
• 首先将输入图像分 割成多个图像片段,然后将它们作为向量线性投射到 Vim 模型中。图像补丁被视为 Vim 块中的序列数据, 这就利用所提出的双向序列状态空间有效地压缩了视觉表示。
• 此外,Vim 块中的位置嵌入提供了对空间信息的感知,这使得 Vim 在密集预测任务中更加稳健。

方法

Vision Mamba (Vim) 的目标是将高级状态空间模型 (SSM)，即 Mamba ，引入计算机视觉。本节首先介绍 SSM 的准备工作。然后概述 Vim。然后，我们详细介绍 Vim 块如何处理输入标记序列，并继续说明 Vim 的架构细节。本节最后分析了所提出的 Vim 的效率。

SSM

略，参加另一篇专门讲Mamba结构的文章

Vision Mamba

图 2 显示了所提出的 Vim 的概览。标准 Mamba 是为 1-D 序列设计的。为了处理视觉任务，我们首先将 2-D 图像 t ∈ R H × W × C t ∈ R^{H×W×C} t∈RH×W×C 转换为扁平的 2-D 块 x p ∈ R J × ( P 2 ⋅ C ) x_p ∈ R^{J×(P^2·C)} xp​∈RJ×(P2⋅C) ，其中 (H, W) 是输入图像的大小，C 是通道数，P 是图像块的大小。接下来，我们将 x p x_p xp​ 线性投影到大小为 D 的向量并添加位置嵌入 E p o s ∈ R ( J + 1 ) × D E_{pos} ∈ R^{(J+1)×D} Epos​∈R(J+1)×D ，如下所示：

其中 t p j t^j_p tpj​ 是 t 的第 j 个 patch， W ∈ R ( P 2 ⋅ C ) × D W ∈ R^{(P^2·C)×D} W∈R(P2⋅C)×D 是可学习的投影矩阵。受 ViT 和 BERT 的启发，我们也使用类 token 来表示整个 patch 序列，记为 t c l s t_{cls} tcls​。然后，我们将 token 序列 $ (T_{l−1})$ 发送到 Vim 编码器的第 l 层，得到输出 Tl。最后，我们对输出类 token T L 0 T^0_L TL0​ 进行归一化，并将其输入到多层感知器 (MLP) 头，得到最终预测 p ^ \hat p p^​，如下所示：

其中 Vim 是提出的视觉曼巴块，L 是层数，Norm 是规范化层。

Vim块

原始 Mamba 块是为 1-D 序列设计的，不适合需要空间感知理解的视觉任务。在本节中，我们介绍了 Vim 块，它结合了视觉任务的双向序列建模。Vim 块如图 2 所示。

具体来说，我们在 Algo.21 中介绍了 Vim 块的操作。输入标记序列 T l − 1 T_{l-1} Tl−1​ 首先由规范化层进行规范化。接下来，我们将规范化的序列线性投影到尺寸大小为 E 的 x 和 z。然后，我们从前向和后向处理 x。对于每个方向，我们首先对 x 应用 1-D 卷积并得到 x ′ o x′o x′o 。然后，我们分别将 x ′ o x′o x′o 线性投影到 B o B_o Bo​、 C o C_o Co​、 ∆ o ∆_o ∆o​。然后分别使用 ∆ o ∆_o ∆o​ 来转换 A o ˉ \bar {A_o} Ao​ˉ​、 B o ˉ \bar {B_o} Bo​ˉ​。最后，我们通过 SSM 计算 y f o r w a r d y_forward yf​orward 和 y b a c k w a r d y_backward yb​ackward 。然后， y f o r w a r d y_forward yf​orward 和 y b a c k w a r d y_backward yb​ackward 由 z 进行门控并相加以获得输出标记序列 T l T_l Tl​。

结构细节

总而言之，我们的架构的超参数如下：

L : t h e n u m b e r o f b l o c k s , L: the number of blocks, L:thenumberofblocks,
D : t h e h i d d e n s t a t e d i m e n s i o n , D: the hidden state dimension, D:thehiddenstatedimension,
E : e x p a n d e d s t a t e d i m e n s i o n , E: expanded state dimension, E:expandedstatedimension,
N : S S M d i m e n s i o n . N: SSM dimension. N:SSMdimension.

参照 ViT 和 DeiT ，我们首先采用 16×16 内核大小的投影层来获得一维非重叠块嵌入序列。随后，我们直接堆叠 L Vim 块。默认情况下，我们将块数 L 设置为 24，将 SSM 维度 N 设置为 16。为了与 DeiT 系列的模型大小保持一致，我们将隐藏状态维度 D 设置为 192，并将扩展状态维度 E 设置为 384（适用于微小尺寸变体）。对于小尺寸变体，我们将 D 设置为 384，将 E 设置为 768。

效率分析

传统的基于 SSM 的方法利用快速傅里叶变换来增强卷积运算，如公式 (4) 所示。对于数据依赖型方法，例如 Mamba，Algo. 21 第 11 行中的 SSM 运算不再等同于卷积。为了解决这个问题，Mamba 和提出的 Vim 选择了一种现代硬件友好的方式来确保效率。这种优化的关键思想是避免现代硬件加速器 (GPU) 的 IO 限制和内存限制。

IO 效率

高带宽内存 (HBM) 和 SRAM 是 GPU 的两个重要组件。其中，SRAM 具有更大的带宽，HBM 具有更大的内存大小。Vim 使用 HBM 的 SSM 操作的标准实现需要大约 O(BMEN) 的内存 IO 数量。受 Mamba 的启发，Vim 首先从慢速 HBM 读入 O ( B M E + E N ) O(BME + EN) O(BME+EN) 字节内存 ( ∆ o , A o , B o , C o ) (∆_o, A_o, B_o, C_o) (∆o​,Ao​,Bo​,Co​) 。然后，Vim 在 SRAM 中获取大小为 (B, M, E, N) 的离散 A o ˉ \bar {A_o} Ao​ˉ​、 B o ˉ \bar {B_o} Bo​ˉ​ 。最后，Vim 在 SRAM 中执行 SSM 操作并将大小为 (B, M, E) 的输出写回 HBM。此方法有助于将 IO 从 O ( B M E N ) O(BMEN) O(BMEN) 减少到 O ( B M E + E N ) O(BME + EN) O(BME+EN)。

内存效率

为了避免内存不足问题并在处理长序列时降低内存使用率，Vim 选择与 Mamba 相同的重新计算方法。对于用于计算梯度的大小为 (B、M、E、N) 的中间状态，Vim 会在网络反向传递时重新计算它们。对于中间激活（例如激活函数和卷积的输出），Vim 也会重新计算它们以优化 GPU 内存需求，因为激活值占用大量内存但重新计算速度很快。

计算效率

Vim 块中的 SSM（Algo.21 中的第 11 行）和 Transformer 中的自注意力机制在自适应提供全局上下文方面都发挥着关键作用。给定一个视觉序列 T ∈ R 1 × M × D T ∈ R^{1×M×D} T∈R1×M×D 和默认设置 E = 2D，全局自注意力机制和 SSM 的计算复杂度为：
其中自注意力与序列长度 M 成二次函数，而 SSM 与序列长度 M 成线性函数（N 是固定参数，默认设置为 16）。计算效率使 Vim 可扩展到具有大序列长度的千兆像素应用。

实验

图像分类

设置

我们在 ImageNet-1K 数据集上对 Vim 进行基准测试，该数据集包含来自 1,000 个类别的 1.28M 张训练图像和 50K 张验证图像。所有模型都在训练集上进行训练，并报告验证集上的 top-1 准确率。为了公平比较，我们的训练设置主要遵循 DeiT。

具体来说，我们应用随机裁剪、随机水平翻转、标签平滑正则化、混合和随机擦除作为数据增强。在对 2242 张输入图像进行训练时，我们使用 AdamW ，动量为 0.9，总批次大小为 1024，权重衰减为 0.05，以优化模型。我们使用余弦计划、 1 × 1 0 − 3 1×10^{−3} 1×10−3 初始学习率和 EMA 对 Vim 模型进行 300 次训练。在测试期间，我们在验证集上应用中心裁剪以裁剪出 22 4 2 224^2 2242 张图像。实验在 8 个 A800 GPU 上进行。长序列微调为了充分利用 Vim 高效的长序列建模能力，我们在 ImageNet 预训练之后继续使用长序列设置对 Vim 进行 30 个 epoch 的微调。具体而言，我们在保持块大小不变的情况下将块提取步长设置为 8，恒定学习率为 1 0 − 5 10^{−5} 10−5，权重衰减为 1 0 − 8 10^{−8} 10−8。

结果

表 1 将 Vim 与基于 ConvNet、Transformer 和 SSM 的骨干网络进行了比较。与基于 ConvNet 的 ResNet 相比，Vim 表现出色。例如，在参数大致相似的情况下，VimSmall 的 top-1 准确率达到 80.5，比 ResNet50 高 4.3 个百分点。与传统的基于自注意力机制的 ViT 相比，Vim 在参数数量和分类准确率方面均有显著提升。与高度优化的 ViT 变体 DeiT 相比，在参数数量相当的情况下，Vim 在不同规模上均超越 DeiT：VimTiny 比 DeiT-Tiny 高 3.9 个百分点，Vim-Small 比 DeiT-Small 高 0.7 个百分点。与基于 SSM 的 S4ND-ViTB 相比，Vim 在参数减少 3 倍的情况下实现了更高的 top-1 准确率。经过长序列微调后，Vim-Tiny† 和 Vim-S† 均取得了更高的结果。其中，VimS† 甚至取得了与 DeiT-B 相似的结果。结果表明，Vim 可以轻松适应更长的序列建模并提取更强的视觉表征。

图 1 (b) 和 © 比较了微小尺寸 Vim 和 DeiT 的 FPS 和 GPU 内存。随着图像分辨率的提高，Vim 在速度和内存方面表现出更好的效率。具体而言，当图像大小为 512×512 时，Vim 实现了与 DeiT 相似的 FPS 和内存。当图像尺寸增长到 1248×1248 时，Vim 比 DeiT 快 2.8 倍，并节省 86.8% 的 GPU 内存。 Vim 在序列长度上的线性缩放的显著优势使其适用于高分辨率下游视觉应用和长序列多模态应用。

语义分割

设置

我们在 ADE20K 上进行语义分割实验，并使用 UperNet 作为分割框架。我们在 B 节结果中提供了详细的设置。如表 2 所示，Vim 在不同规模上的表现始终优于 DeiT：Vim-Ti 比 DeiT-Ti 高 1.8 mIoU，Vim-S 比 DeiT-S 高 0.9 mIoU。与 ResNet-101 主干相比，我们的 Vim-S 以几乎少 2 倍的参数实现了相同的分割性能。为了进一步评估下游任务（即分割、检测和实例分割）的效率，我们将主干与常用的特征金字塔网络 (FPN) 模块相结合，并对它们的 FPS 和 GPU 内存进行基准测试。如图 4 和图 3 所示，效率曲线显示了与纯主干（图 1）类似的比较结果，尽管我们在主干上附加了一个重 FPN。卓越的线性缩放性能归功于我们提出的高效主干 Vim，它为以端到端的方式学习千兆像素级视觉表示奠定了基础，而无需多阶段编码（例如，航空图像、医学图像和计算病理学）。

对象检测和实例分割

设置

我们在 COCO 2017 数据集上进行了对象检测和实例分割实验，并使用 ViTDet 作为基本框架。我们在 B 节结果中提供了详细的设置。表 3 使用级联 Mask R-CNN 框架将 Vim-Ti 与 DeiT-Ti 进行了比较。Vim-Ti 比 DeiTTi 好 1.3 个 box AP 和 1.1 个 mask AP。对于中型和大型物体，Vim-Ti 比 DeiT-Ti 好 1.6 APbox m /1.3 APmask m 和 1.4 APbox l /1.8 APmask l，表现出比 DeiT 更好的远程上下文学习能力（图 5）。我们强调，由于 DeiT 配备了窗口注意功能，而 Vim 以纯序列建模方式工作，因此准确度优势并不简单。具体来说，为了对高分辨率图像（即 1024×1024）进行表征学习，我们遵循 ViTDet 并使用 2D 窗口注意修改 DeiT 主干，这注入了 2D 先验并打破了 Transformer 的顺序建模性质。得益于第 3.5 节图 1 和图 3 中所示的效率，我们可以直接将 Vim 应用于 1024×1024 输入图像并学习顺序视觉表征以进行对象检测和实例分割，而无需在主干中使用 2D 先验。

消融研究

双向 SSM

我们使用 ImageNet-1K 分类和 ADE20K 上的 Segmenter 语义分割框架，消除了 Vim 的关键双向设计。为了充分评估学习表征在 ImageNet 上的能力，我们使用只有 2 层的简单 Segmenter 头对语义分割进行迁移学习。我们研究了以下双向策略：

• 无。我们直接采用 Mamba 块来处理仅具有前向的视觉序列。
• 双向序列。在训练期间，我们随机翻转视觉序列。这就像数据增强一样。
• 双向块。我们将堆叠的块配对。每对的第一个块以前向处理视觉序列，每对的第二个块以后向处理。
• 双向 SSM。我们为每个块添加一个额外的 SSM 来处理后向的视觉序列。
• 双向 SSM + Conv1d。基于双向 SSM，我们进一步在后向 SSM 之前添加一个后向 Conv1d（如图 2）。

如表 4 所示，直接采用 Mamba 块在分类方面取得了良好的效果。然而，不自然的单向方式对下游密集预测提出了挑战。具体来说，使用双向块的初步双向策略在分类上实现了低 7 个点的 top-1 准确率。然而，它在语义分割上的表现比普通单向 Mamba 块高出 1.3 mIoU。通过添加额外的后向 SSM 和 Conv1d，我们实现了卓越的分类准确率（73.9 top-1 acc vs. 73.2 top-1 acc）和卓越的分割优势（35.9 mIoU vs. 32.3 mIoU）。我们在 Vim 块中使用双向 SSM + Conv1d 策略作为默认设置。

分类设计

我们简化了 Vim 的分类设计，以 ImageNet-1K 分类为基准。我们研究了以下分类策略：

• 均值池。我们对最后一个 Vim 块的输出特征采用均值池，并对该池化特征进行分类。
• 最大池。我们首先在视觉序列的每个标记上调整分类头，然后对序列执行最大池以获得分类预测结果。
• 头类标记。按照 DeiT 的方法，我们将类标记连接到视觉序列的头部并进行分类。
• 双类标记。基于头类标记策略，我们在视觉序列的尾部另外添加了一个类标记。
• 中间类标记。我们在视觉序列的中间添加一个类标记，然后对最终的中间类标记进行分类。

如表 5 所示，实验表明，中间类标记策略可以充分利用 SSM 的循环性质和 ImageNet 中的中心对象先验，显示出 76.1 的最佳 top-1 准确率。

总结和展望

• 我们提出了 Vision Mamba (Vim) 来探索最新的高效状态空间模型 Mamba 作为通用视觉主干。
• 与之前使用混合架构或等效全局 2D 卷积核的视觉任务状态空间模型不同，Vim 以序列建模的方式学习视觉表征，并且不会引入特定于图像的归纳偏差。
• 得益于提出的双向状态空间建模，Vim 实现了数据相关的全局视觉上下文，并享有与 Transformer 相同的建模能力，同时计算复杂度更低。
• 受益于 Mamba 的硬件感知设计，Vim 在处理高分辨率图像时，推理速度和内存使用率明显优于 ViT。
• 在标准计算机视觉基准上的实验结果验证了 Vim 的建模能力和高效率，表明 Vim 有很大潜力成为下一代视觉baseline。

在未来的工作中，具有位置嵌入的双向 SSM 建模的 Vim 适用于无监督任务，例如掩模图像建模预训练，而与 Mamba 类似的架构可实现多模态任务，例如 CLIP 样式预训练。基于预训练的 Vim 权重，探索 Vim 在分析高分辨率医学图像、遥感图像和长视频（可视为下游任务）方面的实用性非常简单。